Code import - branch 0.5.0
This commit is contained in:
0
laborious/activities/__init__.py
Normal file
0
laborious/activities/__init__.py
Normal file
169
laborious/activities/activities.py
Normal file
169
laborious/activities/activities.py
Normal file
@@ -0,0 +1,169 @@
|
||||
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 sientia_do.repository.minio_repository import MinioRepository
|
||||
|
||||
from laborious.activities.api import API
|
||||
from laborious.activities.gates import Gates
|
||||
from laborious.activities.mlflow import MLFlow
|
||||
from laborious.activities.model_metrics import ModelMetrics
|
||||
from laborious.activities.opc import OPC
|
||||
from laborious.activities.storage import Storage
|
||||
|
||||
|
||||
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
|
||||
"""
|
||||
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:
|
||||
- Storage: 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
|
||||
- ModelMetrics: Model performance metrics and drift detection
|
||||
- API: PI Web API export operations for industrial systems
|
||||
|
||||
Attributes:
|
||||
postgres_config (dict): PostgreSQL connection configuration
|
||||
mlflow_config (dict): MLFlow server configuration
|
||||
opc_config (dict): OPC server configuration
|
||||
pi_web_api_config (dict): PI Web API 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],
|
||||
pi_web_api_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
|
||||
pi_web_api_config: PI Web API server configuration dictionary
|
||||
Required keys: base_url, auth_type, auth_token
|
||||
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)
|
||||
|
||||
minio_repository = MinioRepository(
|
||||
endpoint=minio_config['endpoint_url'],
|
||||
access_key=minio_config['access_key'],
|
||||
secret_key=minio_config['secret_key'],
|
||||
bucket=minio_config['default_bucket'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
secure=minio_config['secure'],
|
||||
)
|
||||
|
||||
# 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'],
|
||||
retention_hours=minio_config['retention_hours'],
|
||||
minio_repository=minio_repository,
|
||||
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_repository=minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
Gates.__init__(
|
||||
self,
|
||||
minio_repository=minio_repository,
|
||||
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,
|
||||
)
|
||||
|
||||
API.__init__(
|
||||
self,
|
||||
base_url=pi_web_api_config['base_url'],
|
||||
auth_type=pi_web_api_config['auth_type'],
|
||||
auth_token=pi_web_api_config['auth_token'],
|
||||
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
|
||||
- PI Web API client connections
|
||||
- MLFlow model repositories
|
||||
- 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)
|
||||
API.close(self)
|
||||
305
laborious/activities/api.py
Normal file
305
laborious/activities/api.py
Normal file
@@ -0,0 +1,305 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import json
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
from pandas import DataFrame
|
||||
from sientia_do.notifications.handlers import 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.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.pi_web_api_client import PIWebAPIClient
|
||||
|
||||
from laborious import metrics
|
||||
|
||||
|
||||
PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13
|
||||
|
||||
|
||||
class API(SientiaMonitoring):
|
||||
"""
|
||||
PI Web API operations for writing prediction data to PI Web API.
|
||||
|
||||
This class provides Temporal activities for interacting with the PI Web API
|
||||
to write prediction and confidence values to industrial systems. It handles
|
||||
error scenarios gracefully by setting error confidence values and sending
|
||||
notifications when write operations fail.
|
||||
|
||||
The class implements comprehensive error handling for both prediction and
|
||||
confidence value writes, ensuring that partial failures are properly
|
||||
reported and handled.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
auth_type: str,
|
||||
auth_token: str,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize API activity with PI Web API client.
|
||||
|
||||
Args:
|
||||
base_url (str): Base URL of the PI Web API server
|
||||
auth_type (str): Authentication type ('basic' or 'bearer')
|
||||
auth_token (str): Authentication token
|
||||
logger (Logger): Logger instance for operation logging
|
||||
notification_handler (NotificationHandler): Handler for system notifications
|
||||
metrics_controller (MetricsController): Controller for metrics collection
|
||||
"""
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
self.pi_web_api_client = PIWebAPIClient(
|
||||
base_url=base_url,
|
||||
auth_config={
|
||||
'type': auth_type,
|
||||
'token': auth_token,
|
||||
},
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
headers_config={
|
||||
'Content-Type': 'application/json',
|
||||
'Accept': 'application/json',
|
||||
'x-requested-with': 'piwebapistreams',
|
||||
'User-Agent': 'Aig-Laborious-Agent/1.0',
|
||||
},
|
||||
)
|
||||
|
||||
def get_pi_web_api_core_labels(
|
||||
self,
|
||||
metadata: dict[str, Any],
|
||||
operation_type: str = 'write_pi_web_api_data',
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Generate core labels for PI Web API metrics.
|
||||
|
||||
PI Web API metrics in laborious use the shared ``CORE_LABELS`` from
|
||||
``sientia_do``, which includes ``operation_type``. For this reason,
|
||||
operation_type must always be present in emitted labels.
|
||||
|
||||
Args:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata used to derive labels.
|
||||
- operation_type (str): Operation type label for metric cardinality.
|
||||
|
||||
Return:
|
||||
dict[str, Any]: Core labels dictionary including operation_type.
|
||||
"""
|
||||
return super().get_core_labels(
|
||||
metadata=metadata,
|
||||
operation_type=operation_type,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the PI Web API client and shutdown monitoring services.
|
||||
|
||||
This method properly closes all connections and resources associated
|
||||
with the PI Web API client and monitoring services.
|
||||
"""
|
||||
self.pi_web_api_client.close()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
async def process_pi_web_api_response(
|
||||
self,
|
||||
response_data: list[dict[str, Any]],
|
||||
tags: dict[str, str],
|
||||
core_labels: dict[str, str],
|
||||
metadata: dict[str, Any],
|
||||
) -> tuple[int, str]:
|
||||
"""
|
||||
Process the response data from PI Web API write operation.
|
||||
|
||||
Validates that all tags were successfully written, emits metrics for each tag
|
||||
(success or error), and returns the appropriate prediction confidence value.
|
||||
Sets error confidence if any tag write fails or if the number of written tags
|
||||
doesn't match the expected count.
|
||||
|
||||
Args:
|
||||
- response_data (dict[str, Any]): The response data from the PI Web API write operation.
|
||||
- tags (dict[str, str]): The tags that were written to the PI Web API.
|
||||
- core_labels (dict[str, str]): The core labels of the workflow execution.
|
||||
- metadata (dict[str, Any]): The metadata of the workflow execution.
|
||||
Returns:
|
||||
int: Prediction confidence value (0 for success, 13 for errors)
|
||||
"""
|
||||
|
||||
# Convert tags from name:webid to webid:name
|
||||
tags = {w: t for t, w in tags.items()}
|
||||
|
||||
tag_names = list[str](tags.values())
|
||||
|
||||
confidence = 0
|
||||
|
||||
message = ''
|
||||
|
||||
# Evaluate response for each tag
|
||||
written_tags = []
|
||||
for item in response_data:
|
||||
web_id = item.get('WebId')
|
||||
if not web_id:
|
||||
self.error('The response did not contain some WebIds', metadata)
|
||||
continue
|
||||
errors = item.get('Errors', [])
|
||||
tag_name = tags.get(web_id)
|
||||
if not tag_name:
|
||||
self.error(
|
||||
f'The response did not contain the tag name for WebId {web_id}', metadata
|
||||
)
|
||||
continue
|
||||
if errors:
|
||||
self.error(
|
||||
f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata
|
||||
)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT,
|
||||
tags={
|
||||
**core_labels,
|
||||
'tag_name': tag_name,
|
||||
},
|
||||
)
|
||||
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
else:
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT,
|
||||
tags={
|
||||
**core_labels,
|
||||
'tag_name': tag_name,
|
||||
},
|
||||
)
|
||||
written_tags.append(tag_name)
|
||||
|
||||
if len(written_tags) != len(tag_names):
|
||||
message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.'
|
||||
|
||||
self.error(
|
||||
f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}',
|
||||
block='write_pi_web_api_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
)
|
||||
confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
|
||||
return confidence, message
|
||||
|
||||
@activity.defn(name='write_pi_web_api_data')
|
||||
async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
"""
|
||||
Write prediction and confidence data to PI Web API.
|
||||
|
||||
Writes prediction values and confidence scores to PI Web API using configured
|
||||
web IDs. Processes responses to validate writes and emit metrics. Handles errors
|
||||
gracefully by setting error confidence values when writes fail and sending
|
||||
notifications for both prediction and confidence write errors.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): The input data containing:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- pi_web_api_output_config (dict[str, Any]): PI Web API configuration with:
|
||||
- endpoint (str): PI Web API endpoint URL
|
||||
- prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions
|
||||
- confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence
|
||||
- data (dict[str, Any]): Prediction data, its a dataframe converted to dict.
|
||||
Returns:
|
||||
dict[Any, Any]: Data dictionary with potentially modified confidence values
|
||||
If prediction write fails, prediction_confidence is set to error value (13)
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
data = DataFrame(input_data['data'])
|
||||
pi_web_api_output_config = input_data['pi_web_api_output_config']
|
||||
|
||||
self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata)
|
||||
|
||||
raw_prediction_tags = pi_web_api_output_config['prediction_tags']
|
||||
raw_confidence_tags = pi_web_api_output_config['confidence_tags']
|
||||
prediction_tags = list[str](raw_prediction_tags.values())
|
||||
confidence_tags = list(raw_confidence_tags.values())
|
||||
|
||||
core_labels = self.get_pi_web_api_core_labels(metadata)
|
||||
|
||||
prediction_value = data.head(1)['prediction'].values[0]
|
||||
confidence_value = data.head(1)['prediction_confidence'].values[0]
|
||||
|
||||
try:
|
||||
prediction_response = await self.pi_web_api_client.write_value(
|
||||
web_ids=prediction_tags,
|
||||
value={
|
||||
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||
'Value': prediction_value,
|
||||
},
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
confidence, message = await self.process_pi_web_api_response(
|
||||
response_data=prediction_response,
|
||||
tags=raw_prediction_tags,
|
||||
core_labels=core_labels,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
# Preserve incoming confidence/comments on successful PI writes.
|
||||
# Only downgrade confidence or override comments when PI response
|
||||
# explicitly reports a problem (e.g. partial write mismatch).
|
||||
if confidence != 0:
|
||||
data['prediction_confidence'] = confidence
|
||||
if message:
|
||||
data['comments'] = message
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
|
||||
message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}',
|
||||
block='write_pi_web_api_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
self.error(trace, metadata)
|
||||
|
||||
data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
|
||||
data['comments'] = str(e)
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
try:
|
||||
confidence_response = await self.pi_web_api_client.write_value(
|
||||
web_ids=confidence_tags,
|
||||
value={
|
||||
'Timestamp': data.head(1)['timestamp'].values[0],
|
||||
'Value': float(confidence_value),
|
||||
},
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
await self.process_pi_web_api_response(
|
||||
response_data=confidence_response,
|
||||
tags=raw_confidence_tags,
|
||||
core_labels=core_labels,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
|
||||
message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}',
|
||||
block='write_pi_web_api_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
return data.to_dict()
|
||||
804
laborious/activities/gates.py
Normal file
804
laborious/activities/gates.py
Normal file
@@ -0,0 +1,804 @@
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Any
|
||||
|
||||
from pandas import DataFrame
|
||||
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.utils.formatters import create_sample_dict
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
from laborious.utils.filters.conditional_filters import (
|
||||
filter_empty_data,
|
||||
filter_specific_variables_null_values,
|
||||
)
|
||||
from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
|
||||
# Strongly-typed filter function signatures
|
||||
InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool]
|
||||
ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
|
||||
|
||||
# Input filter function mappings
|
||||
input_filter_functions: dict[str, InputFilterFunc] = {
|
||||
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
}
|
||||
|
||||
# Confidence mappings kept separate from function maps to avoid Union types
|
||||
input_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
'CONTINUE': 2,
|
||||
'REPEAT': -1,
|
||||
}
|
||||
|
||||
# MLFlow response filter function mappings
|
||||
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
|
||||
'API_ERROR': api_error_filter,
|
||||
}
|
||||
|
||||
mlflow_response_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
'CONTINUE': 10,
|
||||
'REPEAT': -1,
|
||||
}
|
||||
|
||||
# MLFlow content filter function mappings
|
||||
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
|
||||
'NAN_VALUES': nan_values_filter,
|
||||
'EMPTY_DATA': filter_empty_data,
|
||||
}
|
||||
|
||||
mlflow_content_path_confidence: Mapping[str, int] = {
|
||||
'STOP': -1,
|
||||
'CONTINUE': 18,
|
||||
'REPEAT': -1,
|
||||
}
|
||||
|
||||
|
||||
class Gates(MinioManager):
|
||||
"""
|
||||
Data quality gates and filtering activities for the Laborious system.
|
||||
|
||||
This class implements comprehensive data quality validation and filtering
|
||||
mechanisms that can be applied at different stages of the prediction pipeline.
|
||||
It provides configurable filters with policy-based decision making to ensure
|
||||
data integrity and quality throughout the ML workflow.
|
||||
|
||||
The class supports multiple filter types and implements a flexible policy
|
||||
system that can be configured for different validation requirements. Each
|
||||
filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence
|
||||
scores and detailed comments for monitoring and debugging.
|
||||
|
||||
Attributes:
|
||||
input_filter_functions (dict): Mapping of input filter names to functions
|
||||
mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions
|
||||
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize data quality gates with logging and notification capabilities.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
|
||||
Raises:
|
||||
Exception: If BaseActivity initialization fails
|
||||
"""
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the gates activity and clean up resources.
|
||||
"""
|
||||
|
||||
MinioManager.close(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
|
||||
"""
|
||||
Log dataframe content only when row count is below the configured threshold
|
||||
|
||||
Args:
|
||||
- message (str): Base log message to identify the dataframe in logs
|
||||
- data (Any): Dataframe-like payload to be logged
|
||||
- metadata (dict[str, Any]): Workflow metadata for contextual logging
|
||||
"""
|
||||
self.debug(
|
||||
build_dataframe_debug_message(
|
||||
message=message,
|
||||
data=data,
|
||||
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
|
||||
),
|
||||
metadata,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
|
||||
"""
|
||||
Read filter policy/config keys in a case-insensitive way.
|
||||
|
||||
Args:
|
||||
config (dict[str, Any]): Filter configuration dictionary.
|
||||
|
||||
Return:
|
||||
tuple[str, dict[str, Any]]: Parsed policy and config payload.
|
||||
"""
|
||||
normalized = {str(key).upper(): value for key, value in config.items()}
|
||||
policy = normalized['POLICY']
|
||||
filter_config = normalized.get('CONFIG', {})
|
||||
return policy, filter_config
|
||||
|
||||
@activity.defn(name='input_gate')
|
||||
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Apply input data quality filters and validation.
|
||||
|
||||
This activity validates input data quality using configurable filters
|
||||
before proceeding with ML operations. It applies multiple filter types
|
||||
and returns a path decision based on the filter results and configured
|
||||
policies.
|
||||
|
||||
The method implements a comprehensive filtering system that:
|
||||
1. Applies configured filters to input data
|
||||
2. Evaluates filter results against policy configurations
|
||||
3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT)
|
||||
4. Provides confidence scores and detailed comments
|
||||
5. Handles errors gracefully with notification integration
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for input validation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- filters (dict): Filter configuration and policies
|
||||
- data (dict): Input data to validate
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
|
||||
Raises:
|
||||
Exception: If filter execution fails or configuration is invalid
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
self.info('Performing input gate...', metadata)
|
||||
|
||||
filters = input_data['filters']
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
self._debug_dataframe('Input data:', data, metadata)
|
||||
self.debug(f'Filters: {filters}', metadata)
|
||||
|
||||
# Apply each configured filter
|
||||
for fil, config in filters.items():
|
||||
if fil not in input_filter_functions:
|
||||
self.error(f'Filter {fil} not found', metadata)
|
||||
continue
|
||||
policy, filter_config = self._read_filter_entry(config)
|
||||
try:
|
||||
if input_filter_functions[fil](data, filter_config):
|
||||
self.debug(f'Data not passed the input filter {fil}:{config}', metadata)
|
||||
filter_output.append(policy)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'INTPUT_GATE_ERROR__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
block='input_gate',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
self.info(f'Input gate result: {path_flag}', metadata)
|
||||
return path_flag, input_path_confidence[path_flag], 'Input data with bad quality'
|
||||
|
||||
self.info('Nothing was filtered by the input gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
|
||||
@activity.defn(name='mlflow_response_gate')
|
||||
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Validate MLFlow API response quality and integrity.
|
||||
|
||||
This activity validates MLFlow API responses to ensure they meet quality
|
||||
standards before proceeding with further processing. It applies response-specific
|
||||
filters and determines appropriate path decisions based on response quality.
|
||||
|
||||
The method implements response validation that:
|
||||
1. Applies MLFlow response-specific filters
|
||||
2. Evaluates API response quality and integrity
|
||||
3. Determines path decisions based on response validation results
|
||||
4. Provides confidence scores and detailed validation comments
|
||||
5. Handles API errors and response validation failures
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for response validation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- filters (dict): Response filter configuration and policies
|
||||
- data (dict): MLFlow API response data to validate
|
||||
- type (str): Type of MLFlow operation (transform, predict)
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
|
||||
Raises:
|
||||
Exception: If response validation fails or configuration is invalid
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Performing mlflow response gate...', metadata)
|
||||
raw_data = input_data['data']
|
||||
filters = input_data['filters']
|
||||
|
||||
self.debug(
|
||||
f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata
|
||||
)
|
||||
self.debug(f'Filters: {filters}', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(raw_data)
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
gate_type = input_data['type']
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
comments = []
|
||||
|
||||
status = payload.status or {}
|
||||
|
||||
for fil, config in filters.items():
|
||||
if fil not in mlflow_response_filter_functions:
|
||||
continue
|
||||
policy, filter_config = self._read_filter_entry(config)
|
||||
try:
|
||||
if mlflow_response_filter_functions[fil](status, filter_config):
|
||||
filter_output.append(policy)
|
||||
comments.append(status.get('message', 'Unknown MLFlow API error'))
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
|
||||
message=status.get('message', 'Unknown MLFlow API error'),
|
||||
block='mlflow_gate',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=status.get('traceback'),
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
block='mlflow_gate',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
self.info(f'Mlflow response gate result: {path_flag}', metadata)
|
||||
return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments)
|
||||
|
||||
self.info('Nothing was filtered by the mlflow response gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
|
||||
@activity.defn(name='mlflow_content_gate')
|
||||
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
|
||||
"""
|
||||
Validate MLFlow prediction content quality and integrity.
|
||||
|
||||
This activity validates the content of MLFlow predictions to ensure they
|
||||
meet quality standards before export and persistence. It applies content-specific
|
||||
filters and determines appropriate path decisions based on content quality.
|
||||
|
||||
The method implements content validation that:
|
||||
1. Applies MLFlow content-specific filters
|
||||
2. Evaluates prediction content quality and integrity
|
||||
3. Determines path decisions based on content validation results
|
||||
4. Provides confidence scores and detailed validation comments
|
||||
5. Handles content validation failures and quality issues
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for content validation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- filters (dict): Content filter configuration and policies
|
||||
- data (dict): MLFlow prediction content to validate
|
||||
- type (str): Type of MLFlow operation (transform, predict)
|
||||
- path_priority (list[str]): Priority order for path decisions
|
||||
|
||||
Returns:
|
||||
tuple: (path_flag, confidence, comment)
|
||||
- path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None)
|
||||
- confidence (int): Confidence score for the decision
|
||||
- comment (str): Detailed explanation of the decision
|
||||
|
||||
Raises:
|
||||
Exception: If content validation fails or configuration is invalid
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Performing mlflow content gate...', metadata)
|
||||
|
||||
filters = input_data['filters']
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
gate_type = input_data['type']
|
||||
path_priority = input_data['path_priority']
|
||||
|
||||
filter_output = []
|
||||
|
||||
self._debug_dataframe('Input data:', data, metadata)
|
||||
self.debug(f'Filters: \n {filters}', metadata)
|
||||
|
||||
for fil, config in filters.items():
|
||||
if fil not in mlflow_content_filter_functions:
|
||||
continue
|
||||
policy, filter_config = self._read_filter_entry(config)
|
||||
try:
|
||||
if mlflow_content_filter_functions[fil](data, filter_config):
|
||||
filter_output.append(policy)
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
|
||||
message=f'Data not passed the content filter {fil}:{config}',
|
||||
block='mlflow_gate',
|
||||
level=NotificationLevel.WARNING,
|
||||
attachment_content=data.to_string(),
|
||||
)
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
|
||||
message=f'Error in filter {fil}:{config}: \n {e}',
|
||||
block='mlflow_gate',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
|
||||
for path_flag in path_priority:
|
||||
if path_flag in filter_output:
|
||||
self.info(f'Mlflow content gate result: {path_flag}', metadata)
|
||||
return (
|
||||
path_flag,
|
||||
mlflow_content_path_confidence[path_flag],
|
||||
'Transformed data not passed the content filter',
|
||||
)
|
||||
|
||||
self.info('Nothing was filtered by the mlflow content gate', metadata)
|
||||
|
||||
del data
|
||||
|
||||
return None, 0, ''
|
||||
|
||||
def get_prediction_store_policy(
|
||||
self, prediction_store_policy: str, metadata: dict[str, Any]
|
||||
) -> tuple[str, int]:
|
||||
"""
|
||||
Parse and validate prediction store policy configuration.
|
||||
|
||||
This method parses prediction store policy strings in the format 'type:value'
|
||||
and validates them against allowed policy types and values. It provides
|
||||
sensible defaults for invalid configurations and logs policy validation
|
||||
failures for operational monitoring.
|
||||
|
||||
Supported Policy Types:
|
||||
- 'lts': Latest timestamp - sorts data by timestamp descending
|
||||
- 'erl': Earliest timestamp - sorts data by timestamp ascending
|
||||
|
||||
Args:
|
||||
prediction_store_policy (str): Policy string in format 'type:value'
|
||||
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||
|
||||
Returns:
|
||||
tuple[str, int]: (policy_type, policy_value)
|
||||
- policy_type (str): Validated policy type ('lts' or 'erl')
|
||||
- policy_value (int): Number of rows to retain
|
||||
"""
|
||||
policy_elements = prediction_store_policy.split(':')
|
||||
|
||||
if len(policy_elements) < 2:
|
||||
self.error(
|
||||
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
|
||||
metadata,
|
||||
)
|
||||
return 'lts', 1
|
||||
|
||||
policy_type = policy_elements[0]
|
||||
policy_value = policy_elements[1]
|
||||
|
||||
# If the policy_type is not lts or erl, we use the default policy
|
||||
# If the policty_value is not a number or 0, we use the default policy
|
||||
if (
|
||||
policy_type not in ['lts', 'erl']
|
||||
or not policy_value.isdigit()
|
||||
or int(policy_value) == 0
|
||||
):
|
||||
self.error(
|
||||
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
|
||||
metadata,
|
||||
)
|
||||
return 'lts', 1
|
||||
|
||||
return policy_type, int(policy_value)
|
||||
|
||||
@activity.defn(name='format_transformed_data')
|
||||
async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Format transformed data for storage and export operations.
|
||||
|
||||
This method formats transformed data from MLFlow model transformations
|
||||
into a standardized format suitable for database storage. It converts
|
||||
wide-format data (columns as variables) into long-format (melted)
|
||||
with proper timestamp handling and model identification.
|
||||
|
||||
The formatting process includes:
|
||||
1. Converting input data dictionary to DataFrame
|
||||
2. Extracting timestamps from DataFrame index
|
||||
3. Resetting index to create sequential row numbers
|
||||
4. Melting data from wide format to long format (variable-value pairs)
|
||||
5. Adding model_id for data lineage tracking
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict)
|
||||
- model_id (str): Unique identifier for the ML model
|
||||
|
||||
Returns:
|
||||
dict: Formatted data dictionary with keys:
|
||||
- timestamp (dict): Timestamp values indexed by row number
|
||||
- variable (dict): Variable names indexed by row number
|
||||
- value (dict): Variable values indexed by row number
|
||||
- model_id (dict): Model identifiers indexed by row number
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
|
||||
model_id = input_data['model_id']
|
||||
|
||||
self.info('Formatting transformed data...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
data['timestamp'] = data.index
|
||||
data = data.reset_index(drop=True)
|
||||
|
||||
data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
|
||||
data['model_id'] = model_id
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe=data,
|
||||
minio_repo=self.minio_repository,
|
||||
model_name=input_data['model_name'],
|
||||
operation='transform',
|
||||
workflow_metadata=metadata,
|
||||
last_timestamp=payload.last_timestamp,
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
@activity.defn(name='format_prediction')
|
||||
async def format_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Format prediction data according to configured storage policies.
|
||||
|
||||
This method formats prediction data for storage and export operations.
|
||||
It applies timestamp-based sorting policies, adds metadata fields,
|
||||
and ensures data consistency before persistence. The method supports
|
||||
multiple storage policies for flexible data retention strategies.
|
||||
|
||||
If only one row is present, we use the last timestamp as the timestamp
|
||||
|
||||
Storage Policies:
|
||||
- 'lts:N': Latest timestamp - retains N most recent predictions
|
||||
- 'erl:N': Earliest timestamp - retains N oldest predictions
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- data (dict[str, Any]): Raw prediction data to format
|
||||
- timestamp (str): Timestamp of the data
|
||||
- model_id (str): Unique identifier for the ML model
|
||||
- prediction_confidence (float): Confidence score for the prediction
|
||||
- prediction_store_policy (str): Storage policy in format 'type:value'
|
||||
|
||||
Returns:
|
||||
dict: Formatted prediction data ready for storage and export
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
last_timestamp = input_data['timestamp']
|
||||
prediction_store_policy = input_data['prediction_store_policy']
|
||||
self.info('Formatting prediction...', metadata)
|
||||
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
# Create timestamp column from index and reset index
|
||||
data['timestamp'] = data.index
|
||||
data = data.reset_index(drop=True)
|
||||
|
||||
self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
|
||||
self._debug_dataframe('Prediction data:', data, metadata)
|
||||
|
||||
policy_type, policy_value = self.get_prediction_store_policy(
|
||||
prediction_store_policy, metadata
|
||||
)
|
||||
|
||||
# If data has no timestamp, we use the default timestamp and not sort the data
|
||||
self.info(
|
||||
f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata
|
||||
)
|
||||
|
||||
# If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows
|
||||
if policy_type == 'lts':
|
||||
self.debug('Sorting data by timestamp descending', metadata)
|
||||
data = data.sort_values(by='timestamp', ascending=False)
|
||||
# If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows
|
||||
elif policy_type == 'erl':
|
||||
self.debug('Sorting data by timestamp ascending', metadata)
|
||||
data = data.sort_values(by='timestamp', ascending=True)
|
||||
else:
|
||||
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
|
||||
raise ValueError(f'Invalid policy type: {policy_type}')
|
||||
|
||||
int_policy_value = int(policy_value)
|
||||
|
||||
data = data.head(int_policy_value)
|
||||
if int_policy_value == 1:
|
||||
data['timestamp'] = last_timestamp
|
||||
|
||||
data['model_id'] = input_data['model_id']
|
||||
data['prediction_confidence'] = input_data['prediction_confidence']
|
||||
data['prediction_status'] = 'Good'
|
||||
data['comments'] = ''
|
||||
data = data.sort_values(by='timestamp', ascending=False)
|
||||
data = data.reset_index(drop=True)
|
||||
|
||||
self.info(f'Prediction formatted: {len(data)} rows', metadata)
|
||||
self._debug_dataframe('Prediction data:', data, metadata)
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
@activity.defn(name='format_default_prediction')
|
||||
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Create and format default prediction data for error conditions.
|
||||
|
||||
This method generates default prediction data when the main prediction
|
||||
pipeline encounters errors or quality issues. It creates a standardized
|
||||
data structure with zero values for predictions and useful metadata
|
||||
for operational monitoring and debugging.
|
||||
|
||||
The default prediction serves as a fallback mechanism to:
|
||||
1. Maintain data pipeline continuity during failures
|
||||
2. Provide operational visibility into prediction quality issues
|
||||
3. Enable downstream systems to handle error conditions gracefully
|
||||
4. Support debugging and troubleshooting efforts
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- timestamp (str): Timestamp for the default prediction
|
||||
- model_id (str): Unique identifier for the ML model
|
||||
- prediction_confidence (float): Confidence score (typically low for errors)
|
||||
- comment (str): Error description or operational comment
|
||||
|
||||
Returns:
|
||||
dict: Formatted default prediction data with error indicators
|
||||
"""
|
||||
|
||||
metadata = input_data['metadata']
|
||||
self.debug('Formatting default prediction...', metadata)
|
||||
|
||||
data = DataFrame(
|
||||
{
|
||||
'prediction': [0],
|
||||
'response_time': [0],
|
||||
'timestamp': [input_data['timestamp']],
|
||||
'model_id': [input_data['model_id']],
|
||||
'prediction_confidence': [input_data['prediction_confidence']],
|
||||
'prediction_status': ['Bad'],
|
||||
'comments': [input_data['comment']],
|
||||
}
|
||||
)
|
||||
|
||||
self.info(f'Default prediction formatted: {data.size} rows', metadata)
|
||||
return data.to_dict()
|
||||
|
||||
@activity.defn(name='format_retrain_report')
|
||||
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Format retrain report data for storage and audit trail maintenance.
|
||||
|
||||
This method formats model retraining operation results into a standardized
|
||||
report format suitable for database storage and operational monitoring.
|
||||
It captures retraining status, timestamps, and model version information
|
||||
for comprehensive audit trails and operational visibility.
|
||||
|
||||
The formatting process includes:
|
||||
1. Extracting retraining experiment response data
|
||||
2. Capturing model update report information (version, MLflow IDs)
|
||||
3. Formatting timestamps and status information
|
||||
4. Conditionally including version information for successful retrains
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- experiment_response (dict): Retraining experiment response containing:
|
||||
- success (bool): Retraining operation success status
|
||||
- timestamp (str): Timestamp of the retraining operation
|
||||
- message (str): Status message or error description
|
||||
- update_report (dict): Model update report containing:
|
||||
- version (str): New model version identifier
|
||||
- mlflow_run_id (str): MLflow run identifier
|
||||
- mlflow_experiment_id (str): MLflow experiment identifier
|
||||
- model_id (str): Unique identifier for the ML model
|
||||
- model_name (str): Name of the ML model
|
||||
|
||||
Returns:
|
||||
dict: Formatted retrain report dictionary with keys:
|
||||
- model_id (dict): Model identifiers indexed by row number
|
||||
- model_name (dict): Model names indexed by row number
|
||||
- timestamp (dict): Retraining timestamps indexed by row number
|
||||
- status (dict): Retraining status messages indexed by row number
|
||||
- version (dict, optional): Model versions indexed by row number
|
||||
Only included if experiment_response['success'] is True
|
||||
- mlflow_run_id (dict, optional): MLflow run IDs indexed by row number
|
||||
Only included if experiment_response['success'] is True
|
||||
- mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number
|
||||
Only included if experiment_response['success'] is True
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Formatting retrain report...', metadata)
|
||||
|
||||
experiment_response = input_data['experiment_response']
|
||||
update_report = input_data['update_report']
|
||||
model_id = input_data['model_id']
|
||||
model_name = input_data['model_name']
|
||||
|
||||
report = DataFrame(
|
||||
{
|
||||
'model_id': [model_id],
|
||||
'model_name': [model_name],
|
||||
'timestamp': [experiment_response['timestamp']],
|
||||
'status': [experiment_response['message']],
|
||||
}
|
||||
)
|
||||
|
||||
if experiment_response['success']:
|
||||
# Retrain was successfull
|
||||
report['version'] = update_report['version']
|
||||
report['mlflow_run_id'] = update_report['mlflow_run_id']
|
||||
report['mlflow_experiment_id'] = update_report['mlflow_experiment_id']
|
||||
|
||||
self._debug_dataframe('Retrain report:', report, metadata)
|
||||
|
||||
return report.to_dict()
|
||||
|
||||
@activity.defn(name='write_metrics')
|
||||
async def write_metrics(self, input_data: dict[str, Any]):
|
||||
"""
|
||||
Write prediction performance metrics to Prometheus monitoring system.
|
||||
|
||||
This method records comprehensive metrics for prediction operations,
|
||||
enabling operational monitoring, performance analysis, and alerting.
|
||||
It tracks prediction counts, confidence levels, and response times
|
||||
for each model and pipeline combination.
|
||||
|
||||
Metrics Recorded:
|
||||
1. Prediction Count: Incremental counter for successful predictions
|
||||
2. Confidence Monitor: Current confidence level for predictions
|
||||
3. Response Time Monitor: Histogram of prediction response times
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict[str, Any]): Workflow execution metadata
|
||||
- prediction (dict[str, Any]): Prediction data with metrics
|
||||
|
||||
Raises:
|
||||
Exception: If metrics writing fails or configuration is invalid
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
prediction = DataFrame(input_data['prediction'])
|
||||
prediction_confidence = prediction['prediction_confidence'].values[0]
|
||||
response_time = prediction['response_time'].values[0]
|
||||
opc_metrics = input_data['opc_metrics']
|
||||
|
||||
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
|
||||
|
||||
core_tags = {
|
||||
'pod_id': self.pod_id,
|
||||
'runtime': self.runtime,
|
||||
'operation_type': 'predict',
|
||||
'model_name': metadata['model_name'],
|
||||
'workflow_name': metadata['workflow_name'],
|
||||
}
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
|
||||
tags=core_tags,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
|
||||
method='set',
|
||||
tags=core_tags,
|
||||
value=prediction_confidence,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
|
||||
method='observe',
|
||||
tags=core_tags,
|
||||
value=response_time,
|
||||
)
|
||||
|
||||
for server_id, tags in opc_metrics.items():
|
||||
for tag, response_time in tags.items():
|
||||
if response_time is not None:
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
|
||||
method='observe',
|
||||
tags={
|
||||
**core_tags,
|
||||
'opc_server_id': server_id,
|
||||
'tag': tag,
|
||||
},
|
||||
value=response_time,
|
||||
)
|
||||
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
|
||||
tags={
|
||||
**core_tags,
|
||||
'opc_server_id': server_id,
|
||||
'tag': tag,
|
||||
},
|
||||
)
|
||||
|
||||
self.info(f'Metrics written for model {metadata["model_name"]}', metadata)
|
||||
528
laborious/activities/mlflow.py
Normal file
528
laborious/activities/mlflow.py
Normal file
@@ -0,0 +1,528 @@
|
||||
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.dataframe_debug import build_dataframe_debug_message
|
||||
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
|
||||
"""
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
|
||||
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()
|
||||
|
||||
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
|
||||
"""
|
||||
Log dataframe content only when row count is below the configured threshold
|
||||
|
||||
Args:
|
||||
- message (str): Base log message to identify the dataframe in logs
|
||||
- data (Any): Dataframe-like object expected to expose shape and to_csv
|
||||
- metadata (dict[str, Any]): Workflow metadata for contextual logging
|
||||
"""
|
||||
self.debug(
|
||||
build_dataframe_debug_message(
|
||||
message=message,
|
||||
data=data,
|
||||
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
|
||||
),
|
||||
metadata,
|
||||
)
|
||||
|
||||
@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_dataframe('Raw input data:', data, 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_dataframe('Processed input data:', data, 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,
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
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,
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
@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_dataframe('Input data for prediction:', data, 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,
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
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,
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
@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')
|
||||
364
laborious/activities/model_metrics.py
Normal file
364
laborious/activities/model_metrics.py
Normal file
@@ -0,0 +1,364 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import time
|
||||
import traceback
|
||||
import warnings
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from pandas import DataFrame, Index, to_datetime
|
||||
from sientia.ModelAnalysis import ModelAnalysis
|
||||
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.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from laborious import metrics
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
|
||||
warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0')
|
||||
warnings.filterwarnings(
|
||||
'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide'
|
||||
)
|
||||
|
||||
|
||||
class ModelMetrics(SientiaMonitoring):
|
||||
"""
|
||||
Metrics activities for the Laborious system.
|
||||
|
||||
This class provides activities for writing metrics to the Prometheus monitoring system.
|
||||
"""
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the model metrics activity and clean up resources.
|
||||
"""
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
|
||||
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
|
||||
"""
|
||||
Log dataframe content only when row count is below the configured threshold
|
||||
|
||||
Args:
|
||||
- message (str): Base log message to identify the dataframe in logs
|
||||
- data (Any): Dataframe-like payload to be logged
|
||||
- metadata (dict[str, Any]): Workflow metadata for contextual logging
|
||||
"""
|
||||
self.debug(
|
||||
build_dataframe_debug_message(
|
||||
message=message,
|
||||
data=data,
|
||||
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
|
||||
),
|
||||
metadata,
|
||||
)
|
||||
|
||||
async def get_drift_metrics(
|
||||
self,
|
||||
reference_data: DataFrame,
|
||||
target_data: DataFrame,
|
||||
target_name: str,
|
||||
reference_columns: Index,
|
||||
drift_metrics: list[str],
|
||||
chunk_period: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> DataFrame:
|
||||
"""
|
||||
Calculate univariate drift metrics for a model.
|
||||
Args:
|
||||
model_analysis (ModelAnalysis): Model analysis object
|
||||
reference_data (DataFrame): Reference data
|
||||
target_data (DataFrame): Target data
|
||||
reference_columns (list[str]): Reference columns
|
||||
drift_metrics (list[str]): Drift metrics
|
||||
metadata (dict[str, Any]): Workflow execution metadata
|
||||
"""
|
||||
|
||||
config = {
|
||||
'target': target_name,
|
||||
'prediction': 'prediction',
|
||||
'timestamp': 'timestamp',
|
||||
'features': reference_columns,
|
||||
}
|
||||
|
||||
model_analysis = ModelAnalysis(config=config)
|
||||
|
||||
self._debug_dataframe(
|
||||
f'Reference data: Size {reference_data.shape}', reference_data, metadata
|
||||
)
|
||||
|
||||
self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata)
|
||||
|
||||
core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
|
||||
start_time = time.time()
|
||||
try:
|
||||
univariate_drift = model_analysis.detect_univariate_drift(
|
||||
reference_df=reference_data,
|
||||
analysis_df=target_data,
|
||||
features=reference_columns,
|
||||
timestamp_col=config['timestamp'],
|
||||
methods=drift_metrics,
|
||||
chunk_period=chunk_period,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error detecting univariate drift: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift')
|
||||
start_time = time.time()
|
||||
try:
|
||||
multivariate_drift = model_analysis.detect_multivariate_drift(
|
||||
reference_df=reference_data,
|
||||
analysis_df=target_data,
|
||||
features=reference_columns,
|
||||
timestamp_col=config['timestamp'],
|
||||
chunk_period=chunk_period,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error detecting multivariate drift: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
start_time = time.time()
|
||||
core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe')
|
||||
try:
|
||||
drift_df = model_analysis.get_drift_metrics_dataframe(
|
||||
univariate_drift=univariate_drift,
|
||||
multivariate_drift=multivariate_drift,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||
await self.emit_metric(
|
||||
metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels
|
||||
)
|
||||
raise e
|
||||
await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
|
||||
await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
|
||||
|
||||
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||
|
||||
return drift_df
|
||||
|
||||
@activity.defn(name='calculate_drift')
|
||||
async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
"""
|
||||
Calculate drift metrics for a model.
|
||||
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_name (str): Name of the MLFlow model to calculate drift for
|
||||
- reference_data (pd.DataFrame): Reference data for the model
|
||||
- target_data (pd.DataFrame): Target data for calculating drift
|
||||
- target_name (str): Name of the target column
|
||||
- drift_metrics (list[str]): List of drift metrics to calculate
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
model_name = input_data['model_name']
|
||||
model_id = input_data['model_id']
|
||||
reference_raw_data = input_data['reference_data']
|
||||
target_data = DataFrame(input_data['target_data'])
|
||||
target_name = input_data['target_name']
|
||||
drift_metrics = input_data['drift_metrics']
|
||||
chunk_period = input_data['chunk_period']
|
||||
|
||||
if chunk_period not in ['min', 's']:
|
||||
self.error(f'Invalid chunk period: {chunk_period}', metadata)
|
||||
raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"')
|
||||
|
||||
self.info(f'Calculating drift for model {model_name}', metadata)
|
||||
|
||||
target_data = target_data.pivot(index='timestamp', columns='variable', values='value')
|
||||
target_data['timestamp'] = target_data.index
|
||||
target_data['timestamp'] = to_datetime(target_data['timestamp'])
|
||||
target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
target_data = target_data.reset_index(drop=True)
|
||||
target_data.dropna(inplace=True)
|
||||
|
||||
if reference_raw_data is not None:
|
||||
self.info('Using reference data', metadata)
|
||||
reference_data = DataFrame(reference_raw_data)
|
||||
accurate = True
|
||||
else:
|
||||
# Get 30% first rows of target_data
|
||||
self.warning('Using 30% first rows of target data as reference data', metadata)
|
||||
target_data.sort_values(by='timestamp', ascending=True, inplace=True)
|
||||
reference_data = target_data.head(int(len(target_data) * 0.3))
|
||||
accurate = False
|
||||
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING',
|
||||
message='Using 30% first rows of target data as reference data',
|
||||
block='model_metrics',
|
||||
level=NotificationLevel.WARNING,
|
||||
attachment_content=reference_data.to_csv(),
|
||||
)
|
||||
|
||||
reference_columns = reference_data.drop(
|
||||
columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore'
|
||||
).columns
|
||||
|
||||
try:
|
||||
drift_df = await self.get_drift_metrics(
|
||||
reference_data=reference_data,
|
||||
target_data=target_data,
|
||||
target_name=target_name,
|
||||
reference_columns=reference_columns,
|
||||
drift_metrics=drift_metrics,
|
||||
chunk_period=chunk_period,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
self.error(f'Error getting drift metrics: {e}', metadata)
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR',
|
||||
message=f'Error getting drift metrics: {e}',
|
||||
block='model_metrics',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=traceback.format_exc(),
|
||||
)
|
||||
return []
|
||||
|
||||
if drift_df.empty:
|
||||
self.warning('No drift metrics found', metadata)
|
||||
return []
|
||||
|
||||
# Drop unnecessary columns
|
||||
drift_df.drop(columns=['p_value'], inplace=True)
|
||||
|
||||
# Extract timestamps only until minutes
|
||||
if chunk_period == 'min':
|
||||
target_timestamps = target_data['timestamp'].apply(lambda x: x[:16])
|
||||
else:
|
||||
target_timestamps = target_data['timestamp']
|
||||
|
||||
# Drop rows where timestamp is not in target data, to avoid save drift from reference
|
||||
drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)]
|
||||
|
||||
if drift_df.empty:
|
||||
self.warning(
|
||||
'No drift metrics found after dropping rows where timestamp is not in target data',
|
||||
metadata,
|
||||
)
|
||||
return []
|
||||
|
||||
# Rename columns to match database columns
|
||||
drift_df.rename(
|
||||
columns={
|
||||
'metric': 'method',
|
||||
'statistic': 'value',
|
||||
},
|
||||
inplace=True,
|
||||
)
|
||||
|
||||
# Drop duplicates
|
||||
drift_df.drop_duplicates(
|
||||
subset=['timestamp', 'method', 'feature'], keep='first', inplace=True
|
||||
)
|
||||
|
||||
drift_df['model_id'] = model_id
|
||||
drift_df['accurate'] = accurate
|
||||
|
||||
drift_df['timestamp'] = to_datetime(drift_df['timestamp'])
|
||||
drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC')
|
||||
drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ)
|
||||
|
||||
self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
|
||||
|
||||
return drift_df.to_dict(orient='records')
|
||||
|
||||
@activity.defn(name='calculate_simple_metrics')
|
||||
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]:
|
||||
"""
|
||||
Calculate simple metrics for a model. Metrics available are:
|
||||
- rmse
|
||||
- mse
|
||||
- mae
|
||||
- r2
|
||||
- accuracy
|
||||
- precision
|
||||
- recall
|
||||
- f1
|
||||
Args:
|
||||
input_data (dict[str, Any]): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_id (str): ID of the MLFlow model
|
||||
- target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns
|
||||
- metrics (list[str]): List of metrics to calculate
|
||||
Returns:
|
||||
dict[Hashable, Any]: Dictionary containing the calculated metrics
|
||||
"""
|
||||
|
||||
metadata = input_data['metadata']
|
||||
model_id = input_data['model_id']
|
||||
target_data = DataFrame(input_data['target_data'])
|
||||
metrics = input_data['metrics']
|
||||
interval_minutes = input_data['interval_minutes']
|
||||
|
||||
data_size = target_data.shape[0]
|
||||
|
||||
output_data = []
|
||||
|
||||
diff = target_data['target'] - target_data['prediction']
|
||||
diff_squared = diff**2
|
||||
|
||||
self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata)
|
||||
|
||||
for metric in metrics:
|
||||
if metric == 'rmse':
|
||||
output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))})
|
||||
elif metric == 'mse':
|
||||
output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)})
|
||||
elif metric == 'mae':
|
||||
output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))})
|
||||
elif metric == 'r2':
|
||||
y_true = target_data['target']
|
||||
y_mean = np.mean(y_true)
|
||||
|
||||
ss_res = np.sum(diff_squared)
|
||||
ss_tot = np.sum((y_true - y_mean) ** 2)
|
||||
|
||||
# Evita divisão por zero
|
||||
if ss_tot == 0:
|
||||
r2_score = 0.0
|
||||
else:
|
||||
r2_score = 1 - (ss_res / ss_tot)
|
||||
|
||||
output_data.append({'metric': 'r2', 'value': r2_score})
|
||||
|
||||
data = DataFrame(output_data)
|
||||
data['model_id'] = model_id
|
||||
data['timestamp'] = target_data['timestamp'].max()
|
||||
data['data_size'] = data_size
|
||||
data['interval_minutes'] = interval_minutes
|
||||
|
||||
self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)
|
||||
|
||||
return data.to_dict(orient='records')
|
||||
515
laborious/activities/opc.py
Normal file
515
laborious/activities/opc.py
Normal file
@@ -0,0 +1,515 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from collections.abc import Hashable
|
||||
from typing import Any
|
||||
|
||||
from pandas import DataFrame
|
||||
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.observability.sientia_monitoring import SientiaMonitoring
|
||||
|
||||
from laborious.utils.repository.opc_repository import OpcRepository
|
||||
|
||||
OPC_WRITTING_ERROR_CONFIDENCE = 12
|
||||
OPC_SESSION_BAD_CONFIDENCE = 14
|
||||
OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:'
|
||||
OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers'
|
||||
OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress'
|
||||
OPC_COMMENT_SEPARATOR = ' | '
|
||||
|
||||
|
||||
def _opc_session_bad_comment(opc_status: str | None) -> str:
|
||||
status = opc_status or 'Unknown'
|
||||
return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}'
|
||||
|
||||
|
||||
def _apply_opc_write_error(
|
||||
error_info: dict[str, Any] | None,
|
||||
session_bad_seen: bool,
|
||||
session_bad_status: str | None,
|
||||
reconnect_in_progress_seen: bool,
|
||||
) -> tuple[bool, str | None, bool]:
|
||||
"""
|
||||
Update session/reconnect flags from an OPC write error payload.
|
||||
|
||||
Args:
|
||||
error_info: Repository error details, or None when the write succeeded.
|
||||
session_bad_seen: Whether a session_bad error was seen so far.
|
||||
session_bad_status: Last known OPC status for session errors.
|
||||
reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far.
|
||||
|
||||
Return:
|
||||
Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen).
|
||||
"""
|
||||
if not error_info:
|
||||
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||
|
||||
kind = error_info.get('opc_error_kind')
|
||||
if kind == 'session_bad':
|
||||
return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen
|
||||
if kind == 'reconnect_in_progress':
|
||||
return session_bad_seen, session_bad_status, True
|
||||
return session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||
|
||||
|
||||
class OPC(SientiaMonitoring):
|
||||
"""
|
||||
OPC server integration activities for real-time data export.
|
||||
|
||||
This class provides comprehensive OPC UA client functionality for connecting
|
||||
to multiple OPC servers and writing prediction data in real-time. It implements
|
||||
secure communication with certificate-based authentication and automatic
|
||||
reconnection capabilities.
|
||||
|
||||
The class supports multiple OPC servers with individual configurations and
|
||||
provides robust error handling and monitoring for production environments.
|
||||
|
||||
Attributes:
|
||||
opc_servers (dict): Configuration for multiple OPC servers
|
||||
opc_repository (dict): Active OPC repository connections
|
||||
logger (Logger): Logging and observability instance
|
||||
notification_handler (NotificationHandler): Notification management instance
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
opc_servers: dict[str, dict[str, Any]],
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
self.logger = logger
|
||||
self.notification_handler = notification_handler
|
||||
self.opc_servers = opc_servers
|
||||
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
|
||||
self.opc_repository: dict[str, OpcRepository] = {}
|
||||
|
||||
async def init_opc(self):
|
||||
"""
|
||||
Initialize OPC server connections and establish communication channels.
|
||||
|
||||
This method iterates through all configured OPC servers and attempts to
|
||||
establish secure connections using certificate-based authentication.
|
||||
Each server connection is managed independently, and connection failures
|
||||
are reported through the notification system.
|
||||
|
||||
The method performs the following operations:
|
||||
1. Creates OpcRepository instances for each configured server
|
||||
2. Establishes secure connections with certificate validation
|
||||
3. Reports connection success/failure through notifications
|
||||
4. Logs connection status for operational visibility
|
||||
|
||||
Raises:
|
||||
Exception: If OPC repository initialization fails or connection
|
||||
establishment encounters critical errors
|
||||
|
||||
Note:
|
||||
Connection failures are logged and reported but do not prevent
|
||||
the initialization of other OPC servers. Each server is handled
|
||||
independently to ensure maximum availability.
|
||||
"""
|
||||
self.logger.info('Initializing OPC servers...')
|
||||
for opc_id, server in self.opc_servers.items():
|
||||
self.opc_repository[opc_id] = OpcRepository(
|
||||
opc_id=server['id'],
|
||||
server_name=server['server_name'],
|
||||
url=server['url'],
|
||||
logger=self.logger,
|
||||
server_uri=server['server_uri'],
|
||||
cert_path=server['cert_path'],
|
||||
private_key_path=server['private_key_path'],
|
||||
server_cert_path=server['server_cert_path'],
|
||||
notification_handler=self.notification_handler,
|
||||
reconnection_interval=server['reconnection_interval'],
|
||||
metrics_controller=self.metrics_controller,
|
||||
)
|
||||
is_connected, error_data = await self.opc_repository[opc_id].connect()
|
||||
if not is_connected:
|
||||
await self.send_notification_async(
|
||||
metadata={
|
||||
'model_id': '-',
|
||||
'model_name': '-',
|
||||
'workflow_name': '-',
|
||||
'schedule_name': 'INITIALIZATION',
|
||||
},
|
||||
notification_id=error_data['notification_id'],
|
||||
message=error_data['message'],
|
||||
block=error_data['block'],
|
||||
level=error_data.get('level', NotificationLevel.ERROR),
|
||||
attachment_content=error_data.get('attachment_content', None),
|
||||
)
|
||||
else:
|
||||
self.logger.info(
|
||||
f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
|
||||
)
|
||||
|
||||
async def write_data(
|
||||
self,
|
||||
server_id: str,
|
||||
tag: str,
|
||||
data: Any,
|
||||
data_type: str,
|
||||
tag_type: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> tuple[float | None, dict[str, Any] | None]:
|
||||
"""
|
||||
Write data to a specific OPC server tag with comprehensive error handling.
|
||||
|
||||
Return:
|
||||
tuple[float | None, dict[str, Any] | None]: Response time on success, or
|
||||
(None, error info_data) on repository failure.
|
||||
"""
|
||||
|
||||
try:
|
||||
is_success, info_data = await self.opc_repository[server_id].write_data(
|
||||
tag, data, data_type, metadata
|
||||
)
|
||||
if not is_success:
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=info_data['notification_id'],
|
||||
message=info_data['message'],
|
||||
block=info_data['block'],
|
||||
level=info_data.get('level', NotificationLevel.ERROR),
|
||||
attachment_content=info_data.get('attachment_content', None),
|
||||
)
|
||||
return None, info_data
|
||||
return info_data['response_time'], None
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
|
||||
message=f'Error writing data to OPC server: {e}',
|
||||
block='write_opc_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
raise e
|
||||
|
||||
async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
|
||||
"""
|
||||
Validate that an OPC server is available and configured for write operations.
|
||||
|
||||
This method checks if the specified OPC server exists in the active
|
||||
repository and is available for data writing operations. It provides
|
||||
immediate feedback for server availability and logs validation failures
|
||||
for operational monitoring.
|
||||
|
||||
Args:
|
||||
server_id (str): Unique identifier for the OPC server to validate
|
||||
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||
|
||||
Returns:
|
||||
bool: True if server is available, False otherwise
|
||||
|
||||
Note:
|
||||
Server validation failures are automatically reported through the
|
||||
notification system with detailed information about available servers.
|
||||
This helps operators quickly identify configuration issues.
|
||||
"""
|
||||
if self.opc_repository.get(server_id) is None:
|
||||
message = f'OPC server {server_id} not found to perform write operation.'
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='OPC_SERVER_NOT_FOUND',
|
||||
message=message,
|
||||
block='write_opc_data',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _write_tags_from_config(
|
||||
self,
|
||||
server_id: str,
|
||||
tags_config: dict[str, dict[str, Any]],
|
||||
data: DataFrame,
|
||||
data_column: str,
|
||||
tag_type: str,
|
||||
log_label: str,
|
||||
metadata: dict[str, Any],
|
||||
) -> tuple[dict[str, float | None], bool, str | None, bool]:
|
||||
"""
|
||||
Write a group of OPC tags and collect response times and error flags.
|
||||
|
||||
Args:
|
||||
server_id: Target OPC server identifier.
|
||||
tags_config: Tag name to configuration mapping.
|
||||
data: DataFrame with prediction/confidence columns.
|
||||
data_column: Column name whose first row value is written.
|
||||
tag_type: Tag category passed to write_data ('prediction' or 'confidence').
|
||||
log_label: Human-readable label for success logs.
|
||||
metadata: Context metadata for logging and notifications.
|
||||
|
||||
Return:
|
||||
(response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen)
|
||||
"""
|
||||
response_times: dict[str, float | None] = {}
|
||||
session_bad_seen = False
|
||||
session_bad_status: str | None = None
|
||||
reconnect_in_progress_seen = False
|
||||
|
||||
for tag, tag_config in tags_config.items():
|
||||
response_time, error_info = await self.write_data(
|
||||
server_id=server_id,
|
||||
tag=tag,
|
||||
data=data.head(1)[data_column].values[0],
|
||||
data_type=tag_config['data_type'],
|
||||
tag_type=tag_type,
|
||||
metadata=metadata,
|
||||
)
|
||||
session_bad_seen, session_bad_status, reconnect_in_progress_seen = (
|
||||
_apply_opc_write_error(
|
||||
error_info,
|
||||
session_bad_seen,
|
||||
session_bad_status,
|
||||
reconnect_in_progress_seen,
|
||||
)
|
||||
)
|
||||
if response_time is not None:
|
||||
self.info(
|
||||
f'{log_label} written to OPC server {server_id} for tag {tag}.',
|
||||
metadata,
|
||||
)
|
||||
response_times[tag] = response_time
|
||||
|
||||
return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen
|
||||
|
||||
async def manage_output_tags(
|
||||
self,
|
||||
server_id: str,
|
||||
config: dict[str, Any],
|
||||
data: DataFrame,
|
||||
metadata: dict[str, Any],
|
||||
) -> tuple[bool, dict[str, float | None], bool, str | None, bool]:
|
||||
"""
|
||||
Manage the writing of prediction and confidence data to OPC server tags.
|
||||
|
||||
This method orchestrates the writing of multiple data types to OPC servers
|
||||
based on configuration. It handles both prediction data and confidence
|
||||
values independently, allowing for flexible tag configuration and
|
||||
comprehensive error handling.
|
||||
|
||||
The method supports two main tag types:
|
||||
1. Prediction tags: Write actual prediction values to configured OPC tags
|
||||
2. Confidence tags: Write confidence scores to separate OPC tags
|
||||
|
||||
Args:
|
||||
server_id (str): Unique identifier for the target OPC server
|
||||
config (dict[str, Any]): OPC tag configuration containing:
|
||||
- prediction_tags (dict, optional): Prediction tag configurations
|
||||
- confidence_tags (dict, optional): Confidence tag configurations
|
||||
data (DataFrame): DataFrame containing prediction and confidence data
|
||||
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||
success (bool): Current success status to maintain across operations
|
||||
|
||||
Returns:
|
||||
tuple[bool, int]: (overall_success, total_tags_written)
|
||||
- overall_success: True if all configured tags were written successfully
|
||||
- total_tags_written: Count of successfully written tags
|
||||
"""
|
||||
response_times: dict[str, float | None] = {}
|
||||
session_bad_seen = False
|
||||
session_bad_status: str | None = None
|
||||
reconnect_in_progress_seen = False
|
||||
|
||||
tag_groups = (
|
||||
('prediction_tags', 'prediction', 'prediction', 'Prediction data'),
|
||||
('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'),
|
||||
)
|
||||
for config_key, data_column, tag_type, log_label in tag_groups:
|
||||
if config_key not in config:
|
||||
continue
|
||||
(
|
||||
group_times,
|
||||
group_session_bad,
|
||||
group_status,
|
||||
group_reconnect,
|
||||
) = await self._write_tags_from_config(
|
||||
server_id=server_id,
|
||||
tags_config=config[config_key],
|
||||
data=data,
|
||||
data_column=data_column,
|
||||
tag_type=tag_type,
|
||||
log_label=log_label,
|
||||
metadata=metadata,
|
||||
)
|
||||
response_times.update(group_times)
|
||||
if group_session_bad:
|
||||
session_bad_seen = True
|
||||
session_bad_status = group_status or session_bad_status
|
||||
if group_reconnect:
|
||||
reconnect_in_progress_seen = True
|
||||
|
||||
success = None not in response_times.values()
|
||||
return (
|
||||
success,
|
||||
response_times,
|
||||
session_bad_seen,
|
||||
session_bad_status,
|
||||
reconnect_in_progress_seen,
|
||||
)
|
||||
|
||||
@activity.defn(name='write_opc_data')
|
||||
async def write_opc_data(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
|
||||
"""
|
||||
Write prediction and confidence data to OPC servers. The two writing
|
||||
operations are optional and independent of each other.
|
||||
|
||||
Args:
|
||||
- input_data(dict[str, Any]): The input data. Contains the following keys:
|
||||
- data(dict[str, Any]): The dataframe that contains the data to write
|
||||
to the OPC servers.
|
||||
- opc_output_config(dict[str, Any]): The OPC writing configuration.
|
||||
The keys are the OPC server names and the values contain:
|
||||
- prediction_tags(dict[str, Any]): The tags to write to the OPC servers.
|
||||
- confidence_tags(dict[str, Any]): The tags to write to the OPC servers.
|
||||
|
||||
Returns:
|
||||
- dict[Any, Any]: The data that was written to the OPC servers.
|
||||
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Writing data to OPC servers...', metadata)
|
||||
data = DataFrame(input_data['data'])
|
||||
opc_output_config = input_data['opc_output_config']
|
||||
self.info(f'Data to write: {data.size} rows', metadata)
|
||||
|
||||
success = True
|
||||
session_bad_seen = False
|
||||
session_bad_status: str | None = None
|
||||
reconnect_in_progress_seen = False
|
||||
|
||||
metrics: dict[str, dict[str, float | None]] = {}
|
||||
|
||||
for server_id, config in opc_output_config.items():
|
||||
if not await self.validate_server(server_id, metadata):
|
||||
success = False
|
||||
continue
|
||||
|
||||
(
|
||||
local_success,
|
||||
local_response_times,
|
||||
local_session_bad,
|
||||
local_status,
|
||||
local_reconnect_in_progress,
|
||||
) = await self.manage_output_tags(server_id, config, data, metadata)
|
||||
metrics[server_id] = local_response_times
|
||||
local_count = len(local_response_times)
|
||||
success = success and local_success
|
||||
if local_session_bad:
|
||||
session_bad_seen = True
|
||||
session_bad_status = local_status or session_bad_status
|
||||
if local_reconnect_in_progress:
|
||||
reconnect_in_progress_seen = True
|
||||
|
||||
self.info(
|
||||
f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return (
|
||||
self.process_confidence(
|
||||
data,
|
||||
success,
|
||||
metadata,
|
||||
session_bad=session_bad_seen,
|
||||
opc_status=session_bad_status,
|
||||
reconnect_in_progress=reconnect_in_progress_seen,
|
||||
),
|
||||
metrics,
|
||||
)
|
||||
|
||||
def process_confidence(
|
||||
self,
|
||||
data: DataFrame,
|
||||
success: bool,
|
||||
metadata: dict[str, Any],
|
||||
*,
|
||||
session_bad: bool = False,
|
||||
opc_status: str | None = None,
|
||||
reconnect_in_progress: bool = False,
|
||||
) -> dict[Hashable, Any]:
|
||||
"""
|
||||
Process prediction confidence based on OPC write operation success.
|
||||
|
||||
This method updates the prediction confidence values in the DataFrame
|
||||
based on the success status of OPC server write operations. If any
|
||||
write operations failed, it sets the confidence to a predefined error
|
||||
value to indicate data quality issues.
|
||||
|
||||
The method implements a confidence degradation strategy:
|
||||
- Success: Maintains original confidence values
|
||||
- Failure: Sets confidence to error value for operational awareness
|
||||
|
||||
Args:
|
||||
data (DataFrame): DataFrame containing prediction and confidence data
|
||||
success (bool): Overall success status of OPC write operations
|
||||
metadata (dict[str, Any]): Context metadata for logging and notifications
|
||||
|
||||
Returns:
|
||||
dict[Any, Any]: Processed data as a dictionary with updated confidence values
|
||||
|
||||
Note:
|
||||
The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is
|
||||
used to indicate that data was not successfully exported to OPC servers.
|
||||
This allows downstream systems to handle data quality appropriately.
|
||||
"""
|
||||
|
||||
if not success:
|
||||
comment_parts: list[str] = []
|
||||
confidence = OPC_WRITTING_ERROR_CONFIDENCE
|
||||
|
||||
if session_bad:
|
||||
comment_parts.append(_opc_session_bad_comment(opc_status))
|
||||
confidence = OPC_SESSION_BAD_CONFIDENCE
|
||||
if reconnect_in_progress:
|
||||
comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT)
|
||||
confidence = OPC_SESSION_BAD_CONFIDENCE
|
||||
if not comment_parts:
|
||||
comment_parts.append(OPC_WRITTING_ERROR_MESSAGE)
|
||||
|
||||
comments = OPC_COMMENT_SEPARATOR.join(comment_parts)
|
||||
data['prediction_confidence'] = confidence
|
||||
data['comments'] = comments
|
||||
self.debug(
|
||||
f'OPC write issues, confidence={confidence}, comments={comments}',
|
||||
metadata,
|
||||
)
|
||||
else:
|
||||
self.debug('Data written to OPC servers successfully.', metadata)
|
||||
|
||||
return data.to_dict()
|
||||
|
||||
async def close(self):
|
||||
"""
|
||||
Gracefully shutdown all OPC server connections and cleanup resources.
|
||||
|
||||
This method ensures proper cleanup of all active OPC server connections
|
||||
by calling the disconnect method on each repository instance. It's
|
||||
designed to be called during application shutdown to prevent resource
|
||||
leaks and ensure clean termination.
|
||||
|
||||
The method performs the following cleanup operations:
|
||||
1. Iterates through all active OPC repository connections
|
||||
2. Calls disconnect() on each repository instance
|
||||
3. Allows for graceful connection termination
|
||||
4. Prevents resource leaks and connection hanging
|
||||
|
||||
Note:
|
||||
This method should be called during application shutdown to ensure
|
||||
proper cleanup. It handles all active connections regardless of
|
||||
their current state and provides a clean shutdown experience.
|
||||
"""
|
||||
for opc in self.opc_repository.values():
|
||||
await opc.disconnect()
|
||||
209
laborious/activities/storage.py
Normal file
209
laborious/activities/storage.py
Normal file
@@ -0,0 +1,209 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
# Extend the Temporal Postgres activities for convenient query -> MinIO export
|
||||
import traceback
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
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.activities.postgres import Postgres
|
||||
from sientia_do.temporal.constants import now
|
||||
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
|
||||
_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
|
||||
|
||||
|
||||
class Storage(Postgres, MinioManager):
|
||||
"""
|
||||
Extensions for Postgres activities with a helper to export query results
|
||||
directly to MinIO as Parquet and return the object name.
|
||||
"""
|
||||
|
||||
minio_repository: MinioRepository | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
host: str,
|
||||
port: int,
|
||||
user: str,
|
||||
password: str,
|
||||
dbname: str,
|
||||
min_connections: int,
|
||||
max_connections: int,
|
||||
retention_hours: int = 24,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
self.retention_hours = retention_hours
|
||||
Postgres.__init__(
|
||||
self,
|
||||
host=host,
|
||||
port=port,
|
||||
user=user,
|
||||
password=password,
|
||||
dbname=dbname,
|
||||
min_connections=min_connections,
|
||||
max_connections=max_connections,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
)
|
||||
|
||||
@activity.defn(name='load_query_with_minio_offload')
|
||||
async def load_query_with_minio_offload(
|
||||
self, input_data: dict[str, Any]
|
||||
) -> MinioDataFramePayload:
|
||||
"""
|
||||
Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
|
||||
|
||||
Args (input_data):
|
||||
metadata (dict): Workflow metadata (same as load_custom_query).
|
||||
query (str): SQL query.
|
||||
datetime_columns (list[str], optional): Datetime column names.
|
||||
model_name (str): Model name for object key basename.
|
||||
key_prefix (str, optional): Directory prefix inside the bucket.
|
||||
size_threshold_bytes (int, optional): Override env offload threshold.
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
|
||||
"""
|
||||
if self.minio_repository is None:
|
||||
raise ValueError('Minio repository not initialized')
|
||||
|
||||
metadata: dict = input_data.get('metadata', {})
|
||||
model_name = input_data['model_name']
|
||||
|
||||
rows = await self.load_custom_query(
|
||||
input_data,
|
||||
)
|
||||
if not rows:
|
||||
self.error(
|
||||
'load_query_with_minio_offload failed: No data returned from query', metadata
|
||||
)
|
||||
dataframe = None
|
||||
else:
|
||||
dataframe = pd.DataFrame(rows)
|
||||
|
||||
return await MinioDataFramePayload.from_dataframe(
|
||||
dataframe,
|
||||
minio_repo=self.minio_repository,
|
||||
workflow_metadata=metadata,
|
||||
model_name=model_name,
|
||||
operation='initial',
|
||||
logger=self.logger,
|
||||
)
|
||||
|
||||
@activity.defn(name='export_payload_to_postgres')
|
||||
async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
|
||||
"""
|
||||
Export a payload to PostgreSQL.
|
||||
"""
|
||||
metadata = input_data.get('metadata')
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
return await self.export_data_to_postgres(
|
||||
{
|
||||
**input_data,
|
||||
'data': data,
|
||||
}
|
||||
)
|
||||
|
||||
@activity.defn(name='cleanup_minio_objects_expired')
|
||||
async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Delete objects under the given prefixes that are older than the retention window.
|
||||
|
||||
Args (input_data):
|
||||
metadata (dict): Workflow metadata for logging and metrics.
|
||||
prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
|
||||
"""
|
||||
if self.minio_repository is None:
|
||||
raise ValueError('Minio repository not initialized')
|
||||
|
||||
metadata = input_data.get('metadata', {})
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
prefix = payload.cleanup_prefix()
|
||||
base = now()
|
||||
cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
|
||||
hours=self.retention_hours
|
||||
)
|
||||
|
||||
report: dict[str, Any] = {
|
||||
'failed': {},
|
||||
'deleted': {},
|
||||
'failed_count': 0,
|
||||
'deleted_count': 0,
|
||||
}
|
||||
try:
|
||||
keys = await self.minio_repository.list_objects(
|
||||
prefix=prefix,
|
||||
recursive=True,
|
||||
metadata=metadata,
|
||||
)
|
||||
for key in keys:
|
||||
try:
|
||||
ts = MinioDataFramePayload.parse_object_timestamp(key)
|
||||
if ts is None:
|
||||
continue
|
||||
if ts >= cutoff:
|
||||
continue
|
||||
await self.minio_repository.delete_file(
|
||||
object_name=key,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e:
|
||||
report['failed'][key] = {
|
||||
'success': False,
|
||||
'message': str(e),
|
||||
}
|
||||
report['failed_count'] += 1
|
||||
continue
|
||||
report['deleted'][key] = {
|
||||
'success': True,
|
||||
'message': 'Deleted',
|
||||
}
|
||||
report['deleted_count'] += 1
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
|
||||
message=f'Error cleaning up MinIO objects: {e}',
|
||||
block='cleanup_minio_objects_expired',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
)
|
||||
self.error(trace, metadata)
|
||||
else:
|
||||
# Cleanup success is expected in normal flow; avoid noisy INFO notifications
|
||||
# that do not impact behavior and can flood observability in test runs.
|
||||
self.info('MinIO objects cleaned up successfully', metadata)
|
||||
|
||||
return report
|
||||
|
||||
def close(self) -> None:
|
||||
"""Close Storage resources (MinIO client and Postgres engine)."""
|
||||
Postgres.close(self)
|
||||
MinioManager.close(self)
|
||||
|
||||
def __del__(self):
|
||||
self.close()
|
||||
Reference in New Issue
Block a user