gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1646 Sync laborious/ from local release/SIENTIAPDE-1646 (272e02d)

Browse Source
This commit is contained in:
vitor-aignosi
2026-07-15 17:47:52 -03:00
parent 58f8cb9720
commit c76f24c13a
23 changed files with 4507 additions and 1654 deletions
Download Patch File Download Diff File Expand all files Collapse all files

222
laborious/activities/activities.py
View File

@@ -1,99 +1,177 @@
from temporalio import activity, workflow
from sientia_do.observability.metrics_controller import MetricsController
from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.activities.postgres import Postgres
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger
from laborious.activities.mlflow import MLFlow
from laborious.activities.gates import Gates
from laborious.activities.opc import OPC
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_sync import MinioRepository
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
from sientia_model.model_repository.plugin_store import PluginStore
class Activities(Postgres, MLFlow, Gates, OPC):
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
from laborious.utils.connectors_config import build_mlflow_config
class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
"""
Main activities orchestrator for the Laborious system.
Central orchestrator for all Temporal activities used by Laborious workflows.
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.
Composes Storage (Postgres + MinIO offload), MLFlow (wrapper-based inference and retrain
via ``SientiaMLflowRepository``), Gates (data quality and ML response filters), OPC exports,
drift/simple metrics, and PI Web API writes. The worker constructs one ``Activities`` instance
per process and registers its callables on multiple workers bound to different task queues.
The class implements multiple inheritance to combine specialized functionality:
- Postgres: Database operations and data persistence
- MLFlow: Model inference and transformation operations
- Gates: Data quality validation and filtering mechanisms
- OPC: Real-time data export to OPC servers
MLflow connectivity: unless ``mlflow_repository`` is injected (tests only), this class builds
``SientiaMLflowRepository`` from ``build_mlflow_config()`` so tracking credentials and URL
stay aligned with the rest of Laborious env-based configuration.
Attributes:
postgres_config (dict): PostgreSQL connection configuration
mlflow_config (dict): MLFlow server configuration
opc_config (dict): OPC server configuration
logger (Logger): Logging and observability instance
notification_handler (NotificationHandler): Notification management instance
Inherits and exposes behaviour from mixins; the MLFlow mixin holds ``mlflow_repository``
and ``plugin_store`` after ``__init__``.
"""
def __init__(self,
postgres_config: dict[str, Any],
mlflow_config: dict[str, Any],
opc_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler):
def __init__(
self,
postgres_config: dict[str, Any],
plugin_store: PluginStore,
minio_config: dict[str, Any],
opc_config: dict[str, Any],
pi_web_api_config: dict[str, Any],
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController | None = None,
mlflow_repository: SientiaMLflowRepository | None = None,
):
"""
Initialize the Activities orchestrator with all required configurations.
Wire Postgres, MinIO, MLflow, OPC, gates, metrics, and PI Web API into a single object.
This constructor initializes all parent classes with their respective
configurations and sets up the foundation for all activity operations.
A single ``MetricsController`` instance is created (or reused) and passed to MinIO,
MLflow repository, and all mixins so Prometheus and SDK metrics stay consistent.
Args:
postgres_config: PostgreSQL connection configuration dictionary
Required keys: host, port, user, password, dbname, min_connections, max_connections
mlflow_config: MLFlow server configuration dictionary
Required keys: host, port, username, password
opc_config: OPC server configuration dictionary
Can contain multiple server configurations
logger: Logger instance for observability and debugging
notification_handler: Notification handler for alerts and monitoring
- postgres_config: Host, port, credentials, db name, and pool bounds for Storage.
- plugin_store: ``PluginStore`` instance; the worker must call ``install_runtime`` before
activities run so wrapper code is importable.
- minio_config: Endpoint, keys, bucket, retention, and TLS flag for object storage payloads.
- opc_config: Map of OPC server id to connection settings for ``OPC`` mixin.
- pi_web_api_config: Base URL and auth for ``API`` mixin.
- logger: Structured logger used across all activities.
- notification_handler: Handler for alerts and persisted notifications.
- metrics_controller: Optional shared controller; if ``None``, a new one is created.
- mlflow_repository: Optional ``SientiaMLflowRepository`` for unit/e2e tests; in production
leave unset so the repository is built from environment via ``build_mlflow_config()``.
Raises:
Exception: If any parent class initialization fails
Exception: If any parent ``__init__`` fails (e.g. invalid config keys).
Return:
None
"""
# Initialize parent classes
Postgres.__init__(self, host=postgres_config['host'],
port=postgres_config['port'],
user=postgres_config['user'],
password=postgres_config['password'],
dbname=postgres_config['dbname'],
min_connections=postgres_config['min_connections'],
max_connections=postgres_config['max_connections'],
logger=logger,
notification_handler=notification_handler)
MLFlow.__init__(self, mlflow_host=mlflow_config['host'],
mlflow_port=mlflow_config['port'],
mlflow_username=mlflow_config['username'],
mlflow_password=mlflow_config['password'],
logger=logger,
notification_handler=notification_handler)
mc = metrics_controller or MetricsController(logger=logger)
Gates.__init__(self, logger=logger,
notification_handler=notification_handler)
# Production path: one shared MLflow client for all model registry / tracking calls.
if mlflow_repository is None:
mlflow_cfg = build_mlflow_config()
mlflow_repository = SientiaMLflowRepository(
host=mlflow_cfg['url'],
username=mlflow_cfg['username'],
password=mlflow_cfg['password'],
logger=logger,
notification_handler=notification_handler,
metrics_controller=mc,
)
OPC.__init__(self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler)
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=mc,
secure=minio_config['secure'],
)
async def shutdown(self):
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=mc,
)
MLFlow.__init__(
self,
mlflow_repository=mlflow_repository,
plugin_store=plugin_store,
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mc,
)
Gates.__init__(
self,
minio_repository=minio_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mc,
)
OPC.__init__(
self,
opc_servers=opc_config,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mc,
)
ModelMetrics.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=mc,
)
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=mc,
)
def shutdown(self) -> None:
"""
Gracefully shutdown all activities and clean up resources.
Close database pools, sync clients, and OPC sessions in a defined order.
This method ensures proper cleanup of all resources including:
- PostgreSQL connection pools
- OPC server connections
- Any other resources that need explicit cleanup
Should be invoked on worker exit so connection pools and OPC sessions are released
cleanly before process termination.
The method should be called before the application terminates to ensure
proper resource cleanup and prevent resource leaks.
Return:
None
"""
Postgres.close(self)
await OPC.shutdown(self)
Storage.close(self)
MLFlow.close(self)
Gates.close(self)
OPC.close(self)
ModelMetrics.close(self)
API.close(self)

305
laborious/activities/api.py Normal file
View 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_sync 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)
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
)
self.emit_metric_sync(
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:
self.emit_metric_sync(
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,
)
self.send_notification(
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')
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 = 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 = 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()
self.send_notification(
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 = 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,
)
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()
self.send_notification(
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()

601
laborious/activities/gates.py
View File

@@ -1,57 +1,71 @@
from temporalio import activity, workflow
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.temporal.activities.base import BaseActivity
from sientia_do.observability.logger import Logger
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from sientia_do.formatters import create_sample_dict
from laborious.utils.filters.mlflow_filters import nan_values_filter, api_error_filter
from typing import Any
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository_sync import MinioRepository
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
filter_specific_variables_null_values,
)
from pandas import DataFrame
from laborious import metrics
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 = {
input_filter_functions: dict[str, InputFilterFunc] = {
'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {
'STOP': -1,
'CONTINUE': 2,
'REPEAT': -1
}
}
# 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 = {
mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = {
'API_ERROR': api_error_filter,
'path_confidence': {
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1
},
}
mlflow_response_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 10,
'REPEAT': -1,
}
# MLFlow content filter function mappings
mlflow_content_filter_functions = {
mlflow_content_filter_functions: dict[str, ContentFilterFunc] = {
'NAN_VALUES': nan_values_filter,
'EMPTY_DATA': filter_empty_data,
'path_confidence': {
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1
}
}
mlflow_content_path_confidence: Mapping[str, int] = {
'STOP': -1,
'CONTINUE': 18,
'REPEAT': -1,
}
class Gates(BaseActivity):
class Gates(SientiaMonitoring):
"""
Data quality gates and filtering activities for the Laborious system.
@@ -71,7 +85,16 @@ class Gates(BaseActivity):
mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
"""
def __init__(self, logger: Logger, notification_handler: NotificationHandler):
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.
@@ -82,11 +105,65 @@ class Gates(BaseActivity):
Raises:
Exception: If BaseActivity initialization fails
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
self.minio_repository = minio_repository
SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
@activity.defn(name="input_gate")
async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
def close(self) -> None:
"""
Close the gates activity and clean up resources.
"""
if self.minio_repository is not None:
try:
self.minio_repository.close()
finally:
self.minio_repository = None
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,
)
@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')
def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Apply input data quality filters and validation.
@@ -121,49 +198,52 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Performing input gate...", metadata)
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
path_priority = input_data['path_priority']
filter_output = []
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
self.debug(f"Filters: {filters}", metadata)
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)
self.error(f'Filter {fil} not found', metadata)
continue
policy, filter_config = self._read_filter_entry(config)
try:
if input_filter_functions[fil](data, config['config']):
self.debug(
f"Data not passed the input filter {fil}:{config}", metadata)
filter_output.append(config['policy'])
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()
self.send_notification(
metadata=metadata,
notification_id=f"INTPUT_GATE_ERROR__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="input_gate",
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
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_filter_functions['path_confidence'][path_flag], \
"Input data with bad quality"
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)
return None, 0, ""
self.info('Nothing was filtered by the input gate', metadata)
@activity.defn(name="mlflow_response_gate")
async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
del data
return None, 0, ''
@activity.defn(name='mlflow_response_gate')
def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow API response quality and integrity.
@@ -197,59 +277,67 @@ class Gates(BaseActivity):
Exception: If response validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow response gate...", metadata)
self.info('Performing mlflow response gate...', metadata)
raw_data = input_data['data']
filters = input_data['filters']
data = input_data['data']
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 = payload.retrieve(self.minio_repository, metadata)
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(
f"Input data: \n {create_sample_dict(data, max_items=5, max_depth=2)}", metadata)
self.debug(f"Filters: {filters}", metadata)
comments = []
status = payload.status or {}
for fil, config in filters.items():
if fil not in mlflow_response_filter_functions:
self.error(f"Filter {fil} not found", metadata)
continue
policy, filter_config = self._read_filter_entry(config)
try:
if mlflow_response_filter_functions[fil](data, config):
filter_output.append(config['policy'])
comments.append(data['content']['message'])
if mlflow_response_filter_functions[fil](status, filter_config):
filter_output.append(policy)
comments.append(status.get('message', 'Unknown MLFlow API error'))
self.send_notification(
metadata=metadata,
notification_id=f"{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}",
message=data['content']['message'],
block="mlflow_gate",
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=data['content']['traceback']
attachment_content=status.get('traceback'),
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_RESPONSE_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
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
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_filter_functions['path_confidence'][path_flag], \
", ".join(comments)
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)
return None, 0, ""
self.info('Nothing was filtered by the mlflow response gate', metadata)
@activity.defn(name="mlflow_content_gate")
async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
del data
return None, 0, ''
@activity.defn(name='mlflow_content_gate')
def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
"""
Validate MLFlow prediction content quality and integrity.
@@ -283,56 +371,65 @@ class Gates(BaseActivity):
Exception: If content validation fails or configuration is invalid
"""
metadata = input_data['metadata']
self.info("Performing mlflow content gate...", metadata)
self.info('Performing mlflow content gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
gate_type = input_data['type']
path_priority = input_data['path_priority']
filter_output = []
self.debug(f"Input data:\n {data.head(5).to_string()}", metadata)
self.debug(f"Filters: \n {create_sample_dict(filters)}", metadata)
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, config):
filter_output.append(config['policy'])
if mlflow_content_filter_functions[fil](data, filter_config):
filter_output.append(policy)
self.send_notification(
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",
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()
attachment_content=data.to_string(),
)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"MLFLOW_GATE_CONTENT_FILTER__{fil}",
message=f"Error in filter {fil}:{config}: \n {e}",
block="mlflow_gate",
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
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_filter_functions['path_confidence'][path_flag], \
"Transformed data not passed the content filter"
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)
return None, 0, ""
self.info('Nothing was filtered by the mlflow content gate', metadata)
def get_prediction_store_policy(self,
prediction_store_policy: str,
metadata: dict[str, Any]) -> tuple[str, int]:
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.
@@ -358,7 +455,9 @@ class Gates(BaseActivity):
if len(policy_elements) < 2:
self.error(
f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata)
f'Invalid prediction store policy: {prediction_store_policy}, using default policy',
metadata,
)
return 'lts', 1
policy_type = policy_elements[0]
@@ -366,15 +465,76 @@ class Gates(BaseActivity):
# 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:
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)
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_prediction")
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@activity.defn(name='format_transformed_data')
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 = 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 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')
def format_prediction(self, input_data: dict[str, Any]) -> dict:
"""
Format prediction data according to configured storage policies.
@@ -383,6 +543,8 @@ class Gates(BaseActivity):
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
@@ -390,7 +552,7 @@ class Gates(BaseActivity):
Args:
input_data (dict): Input data containing:
- data (dict[str, Any]): Raw prediction data to format
- timestamp (str): Default timestamp if data lacks timestamp column
- 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'
@@ -399,58 +561,61 @@ class Gates(BaseActivity):
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)
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = 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(f"Prediction data: {data.head(5).to_string()}", metadata)
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)
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)
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)
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)
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}")
self.error(f'Invalid policy type: {policy_type}, using default policy', metadata)
raise ValueError(f'Invalid policy type: {policy_type}')
data = data.head(int(policy_value))
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['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(f"Prediction data: {data.head(5).to_string()}", metadata)
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[Any, Any]:
@activity.defn(name='format_default_prediction')
def format_default_prediction(self, input_data: dict[str, Any]) -> dict:
"""
Create and format default prediction data for error conditions.
@@ -477,65 +642,95 @@ class Gates(BaseActivity):
"""
metadata = input_data['metadata']
self.debug("Formatting default prediction...", 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']]
})
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)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
@activity.defn(name="get_last_timestamp")
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
@activity.defn(name='format_retrain_report')
def format_retrain_report(self, input_data: dict[str, Any]) -> dict:
"""
Extract the most recent timestamp from prediction data.
Format retrain report data for storage and audit trail maintenance.
This method analyzes prediction data to find the latest timestamp,
enabling incremental processing and data continuity tracking.
It handles empty datasets gracefully by returning the current time
as a fallback timestamp.
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 method is essential for:
1. Incremental data processing workflows
2. Data continuity validation
3. Timestamp-based data loading optimization
4. Workflow execution tracking
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:
- data (dict[str, Any]): Prediction data to analyze
- 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:
str: Formatted timestamp string in UTC with timezone
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)
self.info("Getting last timestamp...", 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']
data = DataFrame(input_data['data'])
report = DataFrame(
{
'model_id': [model_id],
'model_name': [model_name],
'timestamp': [experiment_response['timestamp']],
'status': [experiment_response['message']],
}
)
self.debug(f"Input data: {data.head(5).to_string()}", metadata)
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']
if data.empty:
return now().strftime(DATETIME_FORMAT_WITH_TZ)
self._debug_dataframe('Retrain report:', report, metadata)
max_timestamp = max(
data['timestamp'].values.tolist())
return report.to_dict()
self.info(
f"Last timestamp: {max_timestamp}", metadata)
return max_timestamp
@activity.defn(name="write_metrics")
async def write_metrics(self, input_data: dict[str, Any]):
@activity.defn(name='write_metrics')
def write_metrics(self, input_data: dict[str, Any]):
"""
Write prediction performance metrics to Prometheus monitoring system.
@@ -561,27 +756,57 @@ class Gates(BaseActivity):
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)
self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
metrics.PREDICTIONS_WRITTEN_COUNT.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
).inc()
core_tags = {
'pod_id': self.pod_id,
'runtime': self.runtime,
'operation_type': 'predict',
'model_name': metadata['model_name'],
'workflow_name': metadata['workflow_name'],
}
self.emit_metric_sync(
metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
tags=core_tags,
)
metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
).set(prediction_confidence)
self.emit_metric_sync(
metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
method='set',
tags=core_tags,
value=prediction_confidence,
)
metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
pod_id=self.pod_id,
model_name=metadata['model_name'],
pipeline_name=metadata['workflow_name']
).observe(response_time)
self.emit_metric_sync(
metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
method='observe',
tags=core_tags,
value=response_time,
)
self.info(
f"Metrics written for model {metadata['model_name']}", metadata)
for server_id, tags in opc_metrics.items():
for tag, response_time in tags.items():
if response_time is not None:
self.emit_metric_sync(
metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
method='observe',
tags={
**core_tags,
'opc_server_id': server_id,
'tag': tag,
},
value=response_time,
)
self.emit_metric_sync(
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)

848
laborious/activities/mlflow.py
View File

@@ -1,333 +1,649 @@
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
import tempfile
import traceback
from datetime import datetime
from pandas import Timestamp, to_datetime
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ
from sientia_do.temporal.activities.base import BaseActivity
from pathlib import Path
from shutil import rmtree
from typing import Any
import mlflow
import numpy as np
import pandas as pd
from pandas import DataFrame, 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.formatters import create_sample_dict
from laborious.utils.repository.model_repository import MLFlowRepository
from typing import Any
import numpy as np
from pandas import DataFrame
import traceback
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository_sync 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 sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
from sientia_model.model_repository.plugin_store import PluginStore
from laborious.utils.dataframe_debug import build_dataframe_debug_message
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
class MLFlow(BaseActivity):
class MLFlow(SientiaMonitoring):
"""
MLFlow integration activities for model inference operations.
Temporal activities that talk to MLflow through ``SientiaMLflowRepository`` and ``SientiaModel`` wrappers.
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.
Models are resolved by registered name and the ``production`` alias (not by legacy stages or
separate transform/predict flavors). ``get_cached_model`` loads or reuses a wrapper; inference
uses ``wrapper.transform`` / ``wrapper.predict``; retrain uses ``wrapper.retrain`` or
``wrapper.train`` plus ``store_model`` and registry promotion via ``promote_to_alias``.
The class implements comprehensive error handling and logging for all
MLFlow operations, ensuring reliable model inference in production environments.
Large inputs and outputs flow through ``MinioDataFramePayload`` when workflows offload parquet
to MinIO. On failure, transform/predict still return a payload with ``success: False`` and
error details for downstream gates.
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
mlflow_repository: Client for tracking, registry, artifact download, and run lifecycle.
plugin_store: Reference to the store (runtime is installed on the worker; reserved for
future store-backed helpers).
"""
def __init__(self, mlflow_host: str, mlflow_port: int, mlflow_username: str,
mlflow_password: str, logger: Logger, notification_handler: NotificationHandler):
_MAX_DEBUG_DATAFRAME_ROWS = 100
_DEFAULT_MODEL_ALIAS = 'production'
_REFERENCE_ARTIFACT_CANDIDATES = ('evaluation_data.csv', 'test_data.csv')
def __init__(
self,
mlflow_repository: SientiaMLflowRepository,
plugin_store: PluginStore,
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.
Attach shared MLflow and MinIO clients used by all ML activities in this mixin.
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
- mlflow_repository: Repository built by ``Activities`` (or injected in tests).
- plugin_store: Plugin store instance from worker bootstrap.
- minio_repository: MinIO client for ``MinioDataFramePayload`` upload/download.
- logger: Structured logger.
- notification_handler: Notifications on hard failures where applicable.
- metrics_controller: Shared metrics controller.
Raises:
Exception: If MLFlowRepository initialization fails
Return:
None
"""
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
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
self.minio_repository = minio_repository
SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
self.mlflow_repository = mlflow_repository
self.plugin_store = plugin_store
def close(self) -> None:
"""
Release MinIO manager resources held by the mixin.
Return:
None
"""
if self.minio_repository is not None:
try:
self.minio_repository.close()
finally:
self.minio_repository = None
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 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]) -> dict[str, Any]:
def _detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
"""
Transform input data using MLFlow models.
Ensure the transform output index is homogeneous and encoded as ``DATETIME_FORMAT_WITH_TZ`` strings.
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
Accepts an all-string index (validated against the format), or all-``datetime`` /
``Timestamp`` (naive timestamps are localized to UTC before formatting). Mixed element types
or unsupported types raise ``ValueError`` with a message logged at info level.
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
- data: DataFrame whose index carries the time dimension after transform.
- metadata: Workflow metadata for log correlation.
Returns:
dict: Transformed data from MLFlow model
Return:
``pd.DataFrame``: Same frame with a normalized string index; empty frames are returned as-is.
"""
if data.empty:
self.info('Data is empty, skipping datetime index detection and parsing', metadata)
return data
index = data.index
index_type = type(index[0])
self.info(f'Index type: {index_type}', metadata)
message = (
f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, '
f'string in format {DATETIME_FORMAT_WITH_TZ}.'
)
if not all(isinstance(i, index_type) for i in index):
types = map(str, map(type, index))
raise ValueError(f'{message}. Elements are {",".join(types)}')
if index_type is str:
try:
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
except ValueError as e:
raise ValueError(f'{message}. Unable to parse given date format: {e}') from e
elif index_type is datetime or index_type is pd.Timestamp:
idx = data.index
if hasattr(idx, 'tz') and idx.tz is None:
data.index = idx.tz_localize('UTC') # type: ignore[attr-defined]
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined]
else:
raise ValueError(f'{message}. Got {index_type}.')
return data
def _resolve_model_version_for_run(self, run_id: str) -> str:
"""
Map an MLflow ``run_id`` to the latest registered model version that produced that run.
``search_model_versions`` may return multiple versions if the model was registered more than
once for the same run; the highest numeric ``version`` wins so promotion targets the newest
artifact set.
Args:
- run_id: Run UUID from ``retrain_model`` / experiment payload.
Return:
str: Registry version string acceptable by ``promote_to_alias``.
Raises:
Exception: If transformation fails or MLFlow model is unavailable
ValueError: If the filter returns no versions (model not registered for this run).
"""
versions = self.mlflow_repository._client.search_model_versions(
filter_string=f"run_id='{run_id}'"
)
if not versions:
raise ValueError(f'No registered model version found for run_id={run_id}')
latest = max(versions, key=lambda v: int(v.version))
return str(latest.version)
def _resolve_model_alias(self, model_config: dict[str, Any] | None = None) -> str:
"""
Resolve which MLflow alias should be used for model lookup/promotion.
Args:
- model_config: Optional model configuration that may include ``alias``.
Return:
str: Alias name trimmed and normalized; defaults to ``production``.
"""
if not model_config:
return self._DEFAULT_MODEL_ALIAS
alias = str(model_config.get('alias', self._DEFAULT_MODEL_ALIAS)).strip()
return alias or self._DEFAULT_MODEL_ALIAS
@activity.defn(name='request_transform')
def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
"""
Pivot long-format sensor rows, load the production wrapper, and run ``wrapper.transform``.
Expected tabular shape after load: columns including ``variable``, ``timestamp``, ``value``,
and ``created_at`` for deduplication. Data are sorted by ``created_at``, de-duplicated per
``(variable, timestamp)``, pivoted wide, then passed to the model. ``model_config`` may
include ``retention_minutes`` for wrapper cache TTL.
Args:
- input_data: Dict with ``metadata``, ``model_name``, ``data`` (``MinioDataFramePayload``
dict or inline dataframe dict), and optional ``model_config``.
Return:
``MinioDataFramePayload`` with transformed frame and ``success: True``, or a payload
with ``success: False`` and exception details in ``status`` if transform fails.
"""
metadata = input_data['metadata']
self.info('Transforming data...', metadata)
data = DataFrame(input_data['data'])
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
model_alias = self._resolve_model_alias(model_config)
self.debug("Raw input data:", metadata)
self.debug(data.head(5).to_string(), metadata)
self._debug_dataframe('Raw input data:', data, metadata)
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
# Long → wide: keep newest row per (variable, timestamp), then pivot for the wrapper API.
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 = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
# data.reset_index(inplace=True)
data.columns.name = None
data.index.name = None
self.debug("Processed input data:", metadata)
self.debug(data.head(5).to_string(), metadata)
data['timestamp'] = data.index
# Request transformation from MLFlow model
response_data = self.model_monitoring_repository.transform(
model_name, data, model_config, metadata
self._debug_dataframe('Processed input data:', data, metadata)
try:
wrapper = self.mlflow_repository.get_cached_model(
model_name=model_name,
alias=model_alias,
retention_minutes=model_config.get('retention_minutes', 0),
metadata=metadata,
)
transformed_df, transform_meta = wrapper.transform(data)
if transform_meta:
self.info(f'Wrapper transform metadata: {transform_meta}', metadata)
transformed_df = self._detect_and_parse_datetime_index(transformed_df, metadata)
response_data: dict[str, Any] = {'success': True, 'content': transformed_df}
except Exception as e:
response_data = {
'success': False,
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
self.debug(
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.debug(
f"Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Transform 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)
self.info("Data transformed successfully", metadata)
if not response_data.get('success', False):
return 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 response_data
return 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]) -> dict[str, Any]:
@activity.defn(name='request_predict')
def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
"""
Execute predictions using MLFlow models.
Load the production wrapper and call ``wrapper.predict`` on the prepared feature frame.
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
The activity normalizes ``NaN`` to ``None`` for JSON-friendly columns, sets the row index
the same way as ``retrain_model`` (UTC ``DatetimeIndex`` from ``DATETIME_FORMAT_WITH_TZ``),
restores that index on the prediction frame, normalizes the prediction index to
``DATETIME_FORMAT_WITH_TZ`` strings like ``request_transform``, and records ``response_time``.
Non-DataFrame predictions are coerced to a single ``prediction`` column.
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
- input_data: Same envelope as ``request_transform`` (``metadata``, ``model_name``,
``data``, optional ``model_config`` with ``retention_minutes``).
Returns:
dict: Prediction results from MLFlow model
Raises:
Exception: If prediction fails or MLFlow model is unavailable
Return:
``MinioDataFramePayload`` with predictions or error status mirroring transform behaviour.
"""
metadata = input_data['metadata']
self.info('Predicting data...', metadata)
data = DataFrame(input_data['data'])
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
model_alias = self._resolve_model_alias(model_config)
self.debug(f"Input data for: \n {data.head(5).to_string()}", metadata)
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 = self.model_monitoring_repository.predict(
model_name, data, model_config, metadata
data.index = pd.DatetimeIndex(
to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
)
input_index = data.index
try:
wrapper = self.mlflow_repository.get_cached_model(
model_name=model_name,
alias=model_alias,
retention_minutes=model_config.get('retention_minutes', 0),
metadata=metadata,
)
start_time = datetime.now()
predict_data, pred_meta = wrapper.predict({}, data)
end_time = datetime.now()
if pred_meta:
self.info(f'Wrapper predict metadata: {pred_meta}', metadata)
if isinstance(predict_data, DataFrame):
self._debug_dataframe(
'Data received from model prediction:', predict_data, metadata
)
predict_data.columns = pd.Index(['prediction'])
else:
self.debug(
f'Data received from model prediction (not a DataFrame): {predict_data}',
metadata,
)
predict_data = pd.DataFrame(predict_data, columns=['prediction'])
predict_data.index = input_index
predict_data['response_time'] = (end_time - start_time).total_seconds()
predict_data = self._detect_and_parse_datetime_index(predict_data, metadata)
response_data: dict[str, Any] = {'success': True, 'content': predict_data}
except Exception as e:
response_data = {
'success': False,
'content': {'message': str(e), 'traceback': traceback.format_exc()},
}
self.debug(
f"Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}", metadata)
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
metadata,
)
self.info("Data predicted successfully", metadata)
self.info('Data predicted successfully', metadata)
return response_data
if not response_data.get('success', False):
return 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,
)
@activity.defn(name="retrain_model")
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
return 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')
def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
"""
Retrain MLFlow models with updated training data.
Fit an updated wrapper from historical data, then log and register in MLflow.
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
Flow: load long-format data from MinIO → dedupe/pivot like inference prep → require
``model_config['target']`` → read current ``production`` version for ``source_run_id`` tag →
run ``wrapper.retrain`` outside run timing → ``start_run`` with retrain tags → log input
CSV artifact → ``store_model`` and ``log_params``. Does not promote; the workflow calls
``update_production_model`` after validation.
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
- input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
``model_config`` with at least ``target``.
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
Return:
On success: ``success``, ``experiment`` (``run_id``, ``experiment_id``, ``experiment_name``),
``message``, ``timestamp``. On failure: ``success: False``, error fields, and optional trace.
"""
if self.minio_repository is None:
raise ValueError('Minio repository not initialized')
metadata = input_data['metadata']
data = DataFrame(input_data['data'])
try:
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
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)
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')
data = data.pivot(index='timestamp', columns='variable',
values='value')
data.sort_index(inplace=True)
data.reset_index(inplace=True)
data = data.dropna()
data = data.pivot(index='timestamp', columns='variable', values='value')
data.fillna(np.nan, inplace=True)
data.columns.name = None
data.index.name = None
data.index = pd.DatetimeIndex(
to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
)
target = model_config.get('target')
if target is None:
msg = 'model_config must include "target" for retraining'
self.info(msg, metadata)
return {
'success': False,
'experiment': None,
'message': msg,
'traceback': '',
'timestamp': str(timestamp),
}
try:
retrain_output, experiment = self.model_monitoring_repository.retrain_model(
data=data,
model_name=model_name
model_alias = self._resolve_model_alias(model_config)
mv_src = self.mlflow_repository._client.get_model_version_by_alias(
name=model_name,
alias=model_alias,
)
source_run_id = mv_src.run_id
wrapper = self.mlflow_repository.get_cached_model(
model_name=model_name,
alias=model_alias,
retention_minutes=0,
metadata=metadata,
)
# Keep heavy model fitting outside MLflow run timing.
prediction_data = wrapper.retrain(data)
prediction_data.rename(columns={target: 'prediction'}, inplace=True)
# Merge prediction data with retrain data
evaluation_data = pd.merge(
data, prediction_data, left_index=True, right_index=True, how='left'
)
# Rename target column to "target"
evaluation_data.rename(columns={target: 'target'}, inplace=True)
# Reset index and put as column "timestamp"
evaluation_data['timestamp'] = evaluation_data.index
evaluation_data.reset_index(drop=True, inplace=True)
evaluation_data.sort_values(by='timestamp', inplace=True, ascending=True)
run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'
with self.mlflow_repository.start_run(
model_name=model_name,
run_name=run_name,
experiment_name=model_name,
tags={'retrain': 'true', 'source_run_id': source_run_id},
metadata=metadata,
) as run_info:
tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_')
try:
raw_csv = Path(tmp_dir) / 'retrain_input.csv'
evaluation_csv = Path(tmp_dir) / 'evaluation_data.csv'
data.to_csv(raw_csv, index=False)
evaluation_data.to_csv(evaluation_csv, index=False)
mlflow.log_artifact(str(raw_csv))
mlflow.log_artifact(str(evaluation_csv))
finally:
rmtree(tmp_dir, ignore_errors=True)
wrapper.store_model(name=model_name)
self.mlflow_repository.log_params(
{
'retrain': 'true',
'retrain_date': datetime.now().isoformat(),
'source_run_id': source_run_id,
'retrain_samples': str(data.shape),
}
)
experiment_payload = {
'run_id': run_info.run_id,
'experiment_id': run_info.experiment_id,
'experiment_name': model_name,
}
return {
'status': retrain_output,
'timestamp': timestamp,
'experiment': experiment
'success': True,
'experiment': experiment_payload,
'message': 'Model retrained successfully.',
'timestamp': str(timestamp),
}
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='RETRAIN_MODEL_ERROR',
message=f'Error retraining model {model_name}: {e}',
block='retrain_model',
level=NotificationLevel.ERROR,
attachment_content=trace
)
self.error(trace, metadata=metadata)
raise e
error_msg = f'Error retraining model {model_name}: {e}'
self.info(error_msg, metadata)
return {
'success': False,
'experiment': None,
'message': error_msg,
'traceback': traceback.format_exc(),
'timestamp': str(timestamp),
}
@activity.defn(name="update_production_model")
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@activity.defn(name='update_production_model')
def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
"""
Update production model with newly trained model version.
Point the ``production`` alias at the model version registered for the retrain run.
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
Resolves the highest numeric registry version whose ``run_id`` matches
``experiment['run_id']``, then calls ``promote_to_alias``. On failure, sends a notification
and re-raises so the workflow can surface the error.
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
- input_data: ``metadata``, ``model_name``, and ``experiment`` with ``run_id`` and
``experiment_id`` (as returned from ``retrain_model``).
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
Return:
Dict with ``model_name``, promoted ``version``, ``mlflow_run_id``, ``mlflow_experiment_id``.
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
model_id = input_data['model_id']
experiment = input_data['experiment']
timestamp = input_data['timestamp']
status = input_data['status']
self.info(
f'Updating production model {model_name} from experiment {experiment}...', metadata)
f'Updating production model {model_name} from experiment {experiment}...', metadata
)
try:
response = self.model_monitoring_repository.update_production_model(
experiment=experiment,
model_name=model_name
run_id = experiment['run_id']
experiment_id = experiment['experiment_id']
version = self._resolve_model_version_for_run(run_id)
promote_alias = self._resolve_model_alias(input_data.get('model_config'))
self.mlflow_repository.promote_to_alias(
model_name=model_name,
version=version,
alias=promote_alias,
metadata=metadata,
)
report = DataFrame([response])
report['model_id'] = model_id
report['model_name'] = model_name
report['timestamp'] = timestamp
report['status'] = status
self.info(
f'Production model {model_name} updated successfully', metadata)
return report.to_dict()
self.info(f'Production model {model_name} updated successfully', metadata)
return {
'model_name': model_name,
'version': version,
'mlflow_run_id': run_id,
'mlflow_experiment_id': experiment_id,
}
except Exception as e:
trace = traceback.format_exc()
@@ -337,7 +653,103 @@ class MLFlow(BaseActivity):
message=f'Error updating production model {model_name}: {e}',
block='update_production_model',
level=NotificationLevel.ERROR,
attachment_content=trace
attachment_content=trace,
)
self.error(trace, metadata=metadata)
raise e
def _resolve_reference_artifact_name(self, run_id: str) -> str | None:
"""
Pick the first available reference CSV artifact path from the MLflow run.
Candidates are checked in priority order: ``retrain_input.csv``, then ``train_data.csv``.
A path matches when it equals the candidate or ends with ``/<candidate>`` for nested layouts.
Args:
- run_id: MLflow run UUID linked to the production model version.
Return:
Artifact path string for ``download_artifacts``, or ``None`` if no candidate exists.
"""
listed = self.mlflow_repository._client.list_artifacts(run_id)
paths = [file_info.path for file_info in listed]
for candidate in self._REFERENCE_ARTIFACT_CANDIDATES:
for path in paths:
if path == candidate or path.endswith(f'/{candidate}'):
return path
return None
def _find_downloaded_csv(self, tmpdir: str, artifact_name: str) -> Path | None:
"""
Locate a downloaded reference CSV in the temp directory.
Args:
- tmpdir: Directory where ``download_artifacts`` wrote files.
- artifact_name: Basename of the resolved artifact (e.g. ``retrain_input.csv``).
Return:
``Path`` to the CSV file if found, else ``None``.
"""
direct = Path(tmpdir) / artifact_name
if direct.exists():
return direct
matches = list(Path(tmpdir).rglob(artifact_name))
return matches[0] if matches else None
@activity.defn(name='get_reference_data')
def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
"""
Download reference training CSV from the MLflow run linked to the production alias.
Resolves ``retrain_input.csv`` or ``train_data.csv`` via artifact listing before download.
``retrain_input.csv`` is preferred when both exist (most recent retrain snapshot). Used by
drift workflows to compare live data against the reference distribution logged with the model.
Timestamps are normalized to ``DATETIME_FORMAT`` string columns before returning records.
Args:
- input_data: ``metadata``, ``model_name``, and optional ``model_config`` with ``alias``.
Return:
List of row dicts with normalized timestamps, or ``None`` if resolution or load fails.
"""
metadata = input_data['metadata']
model_name = input_data['model_name']
try:
model_alias = self._resolve_model_alias(input_data.get('model_config'))
mv = self.mlflow_repository._client.get_model_version_by_alias(
name=model_name,
alias=model_alias,
)
run_id = mv.run_id
artifact_path = self._resolve_reference_artifact_name(run_id)
if artifact_path is None:
self.warning(f'Reference data not found for model {model_name}', metadata)
return None
artifact_name = Path(artifact_path).name
tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
try:
self.mlflow_repository.download_artifacts(
run_id=run_id,
artifact_path=artifact_path,
dst_path=tmpdir,
metadata=metadata,
)
csv_path = self._find_downloaded_csv(tmpdir, artifact_name)
if csv_path is None:
self.warning(f'Reference data not found for model {model_name}', metadata)
return None
reference_data = pd.read_csv(csv_path)
finally:
rmtree(tmpdir, ignore_errors=True)
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')
except Exception as e:
self.warning(f'Reference data not found for model {model_name}: {e}', metadata)
return None

432
laborious/activities/model_metrics.py Normal file
View File

@@ -0,0 +1,432 @@
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
import pandas as pd
from pandas import DataFrame, Index, Series, 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.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
from sientia_model.analytics.drift_analysis import DriftAnalysis, DriftInsufficientDataError
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 and statistical analysis activities for the Laborious pipeline.
This class centralizes drift/statistical computations and model-quality
aggregates used by scheduled workflows. Besides producing tabular outputs
for persistence, it also emits operational metrics (count, lag, error)
through ``SientiaMonitoring`` so execution health is observable in runtime.
"""
_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:
"""
Shutdown monitoring resources associated with model metrics activities.
This is invoked during worker teardown to flush/close metric controller
internals and prevent dangling telemetry tasks.
"""
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,
)
def _drift_analyze_stage_error(
self,
exc: Exception,
context: str,
metadata: dict[str, Any],
core_labels: dict[str, Any],
) -> None:
"""
Log analyzer failure for a drift stage and increment the analyze error metric.
Args:
- exc (Exception): Failure raised by ``sientia_model``.
- context (str): Short label for the log line (e.g. univariate detection).
- metadata (dict[str, Any]): Workflow metadata for logging.
- core_labels (dict[str, Any]): Tags from ``get_core_labels`` for metrics.
"""
self.error(f'{context}: {exc}', metadata)
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels)
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:
"""
Compute univariate and multivariate drift outputs and merge them into one dataframe.
The method orchestrates three analysis stages (univariate drift,
multivariate drift, and dataframe projection), emitting lag/count/error
metrics for each stage independently so failures are attributable.
Args:
- reference_data (DataFrame): Baseline dataset representing expected behavior.
- target_data (DataFrame): Current analysis dataset to compare against reference.
- target_name (str): Target column name used by ``DriftAnalysis`` config.
- reference_columns (Index): Feature columns evaluated for drift.
- drift_metrics (list[str]): Enabled univariate methods.
- chunk_period (str): Time bucket granularity used by analysis methods.
- metadata (dict[str, Any]): Workflow metadata for logs and notifications.
Return:
DataFrame: Consolidated drift dataframe from ``get_drift_metrics_dataframe`` using
``method`` / ``value`` (and optional ``threshold``, ``drift_type``), ready for
activity-level formatting before Postgres export.
"""
# ``DriftAnalysis`` uses truthiness checks on ``features`` (e.g. ``if not features``);
# a pandas ``Index`` is ambiguous in boolean context — normalize to a list.
feature_names: list[str] = list(reference_columns)
config = {
'target': target_name,
'prediction': 'prediction',
'timestamp': 'timestamp',
'features': feature_names,
}
drift_analysis = DriftAnalysis(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 = drift_analysis.detect_univariate_drift(
reference_df=reference_data,
analysis_df=target_data,
features=feature_names,
timestamp_col=config['timestamp'],
methods=drift_metrics,
chunk_period=chunk_period,
)
except Exception as e:
if isinstance(e, DriftInsufficientDataError):
raise
self._drift_analyze_stage_error(
e, 'Error detecting univariate drift', metadata, core_labels
)
raise
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
self.emit_metric_sync(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 = drift_analysis.detect_multivariate_drift(
reference_df=reference_data,
analysis_df=target_data,
features=feature_names,
timestamp_col=config['timestamp'],
chunk_period=chunk_period,
)
except Exception as e:
if isinstance(e, DriftInsufficientDataError):
raise
self._drift_analyze_stage_error(
e, 'Error detecting multivariate drift', metadata, core_labels
)
raise
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
self.emit_metric_sync(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 = drift_analysis.get_drift_metrics_dataframe(
univariate_drift=univariate_drift,
multivariate_drift=multivariate_drift,
)
except Exception as e:
if isinstance(e, DriftInsufficientDataError):
raise
self._drift_analyze_stage_error(
e, 'Error building drift metrics dataframe', metadata, core_labels
)
raise
self.observe_lag_sync(start_time, metrics.MODEL_ANALYZE_LAG, core_labels)
self.emit_metric_sync(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels)
return drift_df
@staticmethod
def _to_naive_utc(series: Series) -> Series:
"""
Parse ``series`` as datetime and return a TZ-naive UTC copy.
``sientia_model.analytics.drift_analysis.DriftAnalysis`` preserves the
timezone of the input dataframe in its outputs, while target rows
loaded from PostgreSQL come in with ``+00:00``. Forcing both sides of
a comparison to TZ-naive UTC keeps ``isin`` / ``floor`` operations
deterministic regardless of how the analyzer constructs its
timestamps.
Args:
- series (Series): Input series containing datetime-parseable values.
Return:
Series: Datetime64 series with ``tz=None`` representing UTC instants.
"""
parsed = to_datetime(series)
if getattr(parsed.dt, 'tz', None) is not None:
parsed = parsed.dt.tz_convert('UTC').dt.tz_localize(None)
return parsed
@activity.defn(name='calculate_drift')
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
# Keep timestamps as datetime: DriftAnalysis._chunk_dataframe relies on
# ``pd.Grouper(freq=...)`` which rejects string timestamp columns.
target_data['timestamp'] = to_datetime(target_data['timestamp'])
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)
if 'timestamp' in reference_data.columns:
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
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
self.send_notification(
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 = 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:
if isinstance(e, DriftInsufficientDataError):
self.error(str(e), metadata)
notification_id = e.notification_id
notification_message = str(e)
else:
self.error(f'Error getting drift metrics: {e}', metadata)
notification_id = 'MODEL_METRICS_GET_DRIFT_METRICS_ERROR'
notification_message = f'Error getting drift metrics: {e}'
self.send_notification(
metadata=metadata,
notification_id=notification_id,
message=notification_message,
block='model_metrics',
level=NotificationLevel.ERROR,
attachment_content=traceback.format_exc(),
)
raise
# Drop chunks whose floored timestamp does not appear in the analysis window.
# ``DriftAnalysis`` chunks over ``analysis_df``; this only excludes rows that
# do not belong to the current target window (e.g. stray merged reference rows).
target_floor = self._to_naive_utc(target_data['timestamp']).dt.floor(chunk_period)
drift_floor = self._to_naive_utc(drift_df['timestamp']).dt.floor(chunk_period)
drift_df = drift_df[drift_floor.isin(target_floor)]
if drift_df.empty:
self.warning(
'No drift metrics found after dropping rows where timestamp is not in target data',
metadata,
)
return []
# Analyzer emits diagnostic columns that are not stored in ``sientia_data.drift_metrics``.
drift_df = drift_df.drop(columns=['threshold', 'drift_type'], errors='ignore')
drift_df['model_id'] = str(model_id)
drift_df['accurate'] = accurate
# ``timestamp`` is overridden with the most recent target instant so
# every persisted row shares a single business timestamp (the run's
# logical "now"), matching what downstream consumers expect.
latest_target_timestamp = self._to_naive_utc(target_data['timestamp']).max()
drift_df['timestamp'] = (
pd.Timestamp(latest_target_timestamp)
.tz_localize('UTC')
.strftime(DATETIME_FORMAT_WITH_TZ)
)
# ``chunk_start_date`` / ``chunk_end_date`` may carry nanosecond
# precision (beyond ``timestamptz`` microseconds), so serialize as ISO
# text for the ``text`` Postgres columns.
for column in ('chunk_start_date', 'chunk_end_date'):
drift_df[column] = drift_df[column].apply(
lambda value: pd.Timestamp(value).isoformat() if pd.notna(value) else None
)
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')
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'])
metric_names = 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}: {metric_names}', metadata)
for metric in metric_names:
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')

388
laborious/activities/opc.py
View File

@@ -1,20 +1,62 @@
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.temporal.activities.base import BaseActivity
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
from typing import Any
import traceback
from pandas import DataFrame
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 = ' | '
class OPC(BaseActivity):
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.
@@ -33,20 +75,20 @@ class OPC(BaseActivity):
notification_handler (NotificationHandler): Notification management instance
"""
def __init__(self, opc_servers: dict[str, dict[str, Any]],
logger: Logger, notification_handler: NotificationHandler):
self.logger = logger
self.notification_handler = notification_handler
def __init__(
self,
opc_servers: dict[str, dict[str, Any]],
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
):
self.opc_servers = opc_servers
BaseActivity.__init__(
self, logger, notification_handler, set_error_counter=True)
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.opc_repository: dict[str, OpcRepository] = {}
self.opc_servers = opc_servers
async def init_opc(self):
def init_opc(self):
"""
Initialize OPC server connections and establish communication channels.
@@ -70,10 +112,11 @@ class OPC(BaseActivity):
the initialization of other OPC servers. Each server is handled
independently to ensure maximum availability.
"""
self.logger.info("Initializing OPC servers...")
for id, server in self.opc_servers.items():
self.opc_repository[id] = OpcRepository(
id=server['id'],
self.info('Initializing OPC servers...')
for opc_id, server in self.opc_servers.items():
self.opc_repository[opc_id] = OpcRepository(
opc_id=opc_id,
server_name=server['server_name'],
url=server['url'],
logger=self.logger,
server_uri=server['server_uri'],
@@ -81,76 +124,70 @@ class OPC(BaseActivity):
private_key_path=server['private_key_path'],
server_cert_path=server['server_cert_path'],
notification_handler=self.notification_handler,
reconnection_interval=server['reconnection_interval'],
pod_id=self.pod_id
reconnection_interval=server.get('reconnection_interval', 60),
metrics_controller=self.metrics_controller,
)
is_connected, error_data = await self.opc_repository[id].connect()
is_connected, error_data = self.opc_repository[opc_id].connect()
if not is_connected:
self.send_notification(
metadata={
'model_id': '-',
'model_name': '-',
'workflow_name': '-',
'schedule_name': 'INITIALIZATION'
'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)
attachment_content=error_data.get('attachment_content', None),
)
else:
self.logger.info(
f"OPC server {id} connected successfully.")
self.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]) -> bool:
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.
This method provides a secure and reliable way to write data to OPC servers
with automatic error handling, notification integration, and detailed logging.
It validates server availability before attempting write operations and
provides comprehensive error reporting for operational monitoring.
Args:
- server_id (str): The id of the OPC server.
- tag (str): The tag to write to.
- data (Any): The data to write.
- data_type (str): The data type.
- tag_type (str): The tag type.
Returns:
- bool: True if the data was written successfully, False otherwise.
Return:
tuple[float | None, dict[str, Any] | None]: Response time on success, or
(None, error info_data) on repository failure.
"""
try:
is_success, error_data = await self.opc_repository[server_id].write_data(
tag, data, data_type, self.logger, metadata)
is_success, info_data = self.opc_repository[server_id].write_data(
tag, data, data_type, metadata
)
if not is_success:
self.send_notification(
metadata=metadata,
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)
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 False
return True
return None, info_data
return info_data['response_time'], None
except Exception as e:
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id=f"WRITE_OPC_{tag_type.upper()}_ERROR",
message=f"Error writing data to OPC server: {e}",
block="write_opc_data",
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
attachment_content=trace,
)
raise e
raise
def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
"""
@@ -174,21 +211,81 @@ class OPC(BaseActivity):
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."
message = f'OPC server {server_id} not found to perform write operation.'
self.send_notification(
metadata=metadata,
notification_id="OPC_SERVER_NOT_FOUND",
notification_id='OPC_SERVER_NOT_FOUND',
message=message,
block="write_opc_data",
block='write_opc_data',
level=NotificationLevel.ERROR,
attachment_content=f"OPC servers: {list(self.opc_repository.keys())}"
attachment_content=f'OPC servers: {list(self.opc_repository.keys())}',
)
return False
return True
async def manage_output_tags(
self, server_id: str, config: dict[str, Any], data: DataFrame,
metadata: dict[str, Any], success: bool) -> tuple[bool, int]:
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 = 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
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.
@@ -215,44 +312,52 @@ class OPC(BaseActivity):
- 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
count = 0
if 'prediction_tags' in config:
for tag, tag_config in config['prediction_tags'].items():
local_success = await self.write_data(
server_id=server_id,
tag=tag,
data=data.head(1)['prediction'].values[0],
data_type=tag_config['data_type'],
tag_type='prediction',
metadata=metadata
)
if local_success:
self.info(
f"Prediction data written to OPC server {server_id} for tag {tag}.", metadata)
count += 1
success = success and local_success
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,
) = 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
if 'confidence_tags' in config:
for tag, tag_config in config['confidence_tags'].items():
local_success = await self.write_data(
server_id=server_id,
tag=tag,
data=data.head(1)['prediction_confidence'].values[0],
data_type=tag_config['data_type'],
tag_type='confidence',
metadata=metadata
)
if local_success:
self.info(
f"Confidence data written to OPC server {server_id} for tag {tag}.", metadata)
count += 1
success = success and local_success
return success, count
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]) -> dict[Any, Any]:
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.
@@ -271,29 +376,68 @@ class OPC(BaseActivity):
"""
metadata = input_data['metadata']
self.info("Writing data to OPC servers...", 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)
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
opc_metrics: dict[str, dict[str, float | None]] = {}
for server_id, config in opc_output_config.items():
if not self.validate_server(server_id, metadata):
success = False
continue
local_success, local_count = await self.manage_output_tags(
server_id, config, data, metadata, success)
(
local_success,
local_response_times,
local_session_bad,
local_status,
local_reconnect_in_progress,
) = self.manage_output_tags(server_id, config, data, metadata)
opc_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
n_pred = len(config.get('prediction_tags') or {})
n_conf = len(config.get('confidence_tags') or {})
self.info(
f"Process completed for OPC server {server_id}: {local_count} of {len(config['prediction_tags'])} prediction tags and {len(config['confidence_tags'])} confidence tags", metadata)
f'Process completed for OPC server {server_id}: {local_count} of {n_pred} prediction tags and {n_conf} confidence tags',
metadata,
)
return self.process_confidence(data, success, 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,
),
opc_metrics,
)
def process_confidence(self, data: DataFrame, success: bool, metadata: dict[str, Any]) -> dict[Any, Any]:
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.
@@ -321,18 +465,31 @@ class OPC(BaseActivity):
"""
if not success:
data['prediction_confidence'] = OPC_WRITTING_ERROR_CONFIDENCE
self.debug(
f"Some data could not be written to OPC servers, setting confidence to {OPC_WRITTING_ERROR_CONFIDENCE}.",
metadata
)
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)
self.debug('Data written to OPC servers successfully.', metadata)
return data.to_dict()
async def shutdown(self):
def close(self):
"""
Gracefully shutdown all OPC server connections and cleanup resources.
@@ -353,4 +510,5 @@ class OPC(BaseActivity):
their current state and provides a clean shutdown experience.
"""
for opc in self.opc_repository.values():
await opc.disconnect()
opc.disconnect()
self.opc_repository.clear()

231
laborious/activities/storage.py Normal file
View File

@@ -0,0 +1,231 @@
from temporalio import activity, workflow
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.observability.sientia_monitoring import SientiaMonitoring
from sientia_do.repository.minio_repository_sync import MinioRepository
from sientia_do.temporal.activities.postgres_sync 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, SientiaMonitoring):
"""
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,
)
self.minio_repository = minio_repository
SientiaMonitoring.__init__(
self,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
)
@activity.defn(name='load_query_with_minio_offload')
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 = 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 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')
def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
"""
Resolve a MinIO-aware payload into a DataFrame and persist it into PostgreSQL.
This activity accepts the serialized payload produced by previous steps
(inline dict or MinIO object reference), reconstructs the tabular data,
and delegates the final write to ``export_data_to_postgres`` using the
same input contract expected by the Postgres activity mixin.
Args:
- input_data (dict[str, Any]): Activity input containing ``data`` as a
``MinioDataFramePayload``-compatible dict plus database write options
(schema/table/on_conflict/metadata and related fields).
Return:
dict: Result dictionary returned by ``export_data_to_postgres``, including
success status and optional write diagnostics.
"""
metadata = input_data.get('metadata')
payload = MinioDataFramePayload.from_dict(input_data['data'])
data = payload.retrieve(self.minio_repository, metadata)
return self.export_data_to_postgres(
{
**input_data,
'data': data,
}
)
@activity.defn(name='cleanup_minio_objects_expired')
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 = 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
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()
self.send_notification(
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:
"""
Shutdown Storage resources in deterministic order.
The method first closes Postgres resources via ``Postgres.close`` (engine,
sessions, and monitoring hooks), then closes the optional MinIO repository
and clears the local reference to avoid accidental reuse after shutdown.
"""
Postgres.close(self)
if self.minio_repository is not None:
try:
self.minio_repository.close()
finally:
self.minio_repository = None