From 33d6ddc26b102bd564f7e0ab5787ab299c640f9f Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 25 Aug 2026 11:11:31 -0300 Subject: [PATCH] Add laborious --- laborious/__init__.py | 0 laborious/activities/__init__.py | 0 laborious/activities/activities.py | 212 ++ laborious/activities/api.py | 305 +++ laborious/activities/gates.py | 804 ++++++++ laborious/activities/mlflow.py | 528 +++++ laborious/activities/model_import.py | 1044 ++++++++++ laborious/activities/model_import_errors.py | 348 ++++ laborious/activities/model_metrics.py | 364 ++++ laborious/activities/opc.py | 515 +++++ laborious/activities/storage.py | 210 ++ laborious/metrics.py | 239 +++ laborious/utils/__init__.py | 0 laborious/utils/bundle/__init__.py | 27 + laborious/utils/bundle/format.py | 216 ++ laborious/utils/bundle/reader.py | 473 +++++ laborious/utils/bundle/steps.py | 55 + laborious/utils/connectors_config.py | 162 ++ laborious/utils/dataframe_debug.py | 34 + laborious/utils/filters/__init__.py | 0 .../utils/filters/conditional_filters.py | 48 + laborious/utils/filters/mlflow_filters.py | 64 + laborious/utils/import_password.py | 98 + laborious/utils/models/__init__.py | 0 .../utils/models/minio_dataframe_payload.py | 348 ++++ laborious/utils/repository/minio_manager.py | 32 + .../utils/repository/model_repository.py | 1774 +++++++++++++++++ laborious/utils/repository/opc_repository.py | 866 ++++++++ laborious/worker/__init__.py | 0 laborious/worker/worker.py | 336 ++++ laborious/workflows/__init__.py | 0 laborious/workflows/drift.py | 107 + laborious/workflows/import_model.py | 416 ++++ laborious/workflows/minimal_retrain.py | 137 ++ laborious/workflows/predictions_batch.py | 127 ++ laborious/workflows/simple_metrics.py | 95 + laborious/workflows/sub_workflows/__init__.py | 0 .../format_and_export_prediction.py | 220 ++ .../sub_workflows/prediction_process.py | 346 ++++ 39 files changed, 10550 insertions(+) create mode 100644 laborious/__init__.py create mode 100644 laborious/activities/__init__.py create mode 100644 laborious/activities/activities.py create mode 100644 laborious/activities/api.py create mode 100644 laborious/activities/gates.py create mode 100644 laborious/activities/mlflow.py create mode 100644 laborious/activities/model_import.py create mode 100644 laborious/activities/model_import_errors.py create mode 100644 laborious/activities/model_metrics.py create mode 100644 laborious/activities/opc.py create mode 100644 laborious/activities/storage.py create mode 100644 laborious/metrics.py create mode 100644 laborious/utils/__init__.py create mode 100644 laborious/utils/bundle/__init__.py create mode 100644 laborious/utils/bundle/format.py create mode 100644 laborious/utils/bundle/reader.py create mode 100644 laborious/utils/bundle/steps.py create mode 100644 laborious/utils/connectors_config.py create mode 100644 laborious/utils/dataframe_debug.py create mode 100644 laborious/utils/filters/__init__.py create mode 100644 laborious/utils/filters/conditional_filters.py create mode 100644 laborious/utils/filters/mlflow_filters.py create mode 100644 laborious/utils/import_password.py create mode 100644 laborious/utils/models/__init__.py create mode 100644 laborious/utils/models/minio_dataframe_payload.py create mode 100644 laborious/utils/repository/minio_manager.py create mode 100644 laborious/utils/repository/model_repository.py create mode 100644 laborious/utils/repository/opc_repository.py create mode 100644 laborious/worker/__init__.py create mode 100644 laborious/worker/worker.py create mode 100644 laborious/workflows/__init__.py create mode 100644 laborious/workflows/drift.py create mode 100644 laborious/workflows/import_model.py create mode 100644 laborious/workflows/minimal_retrain.py create mode 100644 laborious/workflows/predictions_batch.py create mode 100644 laborious/workflows/simple_metrics.py create mode 100644 laborious/workflows/sub_workflows/__init__.py create mode 100644 laborious/workflows/sub_workflows/format_and_export_prediction.py create mode 100644 laborious/workflows/sub_workflows/prediction_process.py diff --git a/laborious/__init__.py b/laborious/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/activities/__init__.py b/laborious/activities/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py new file mode 100644 index 0000000..f4a9a2c --- /dev/null +++ b/laborious/activities/activities.py @@ -0,0 +1,212 @@ +from sientia_do.observability.metrics_controller import MetricsController +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any + + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import Logger + from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.repository.mongodb_repository import MongoDBRepository + + from laborious.activities.api import API + from laborious.activities.gates import Gates + from laborious.activities.mlflow import MLFlow + from laborious.activities.model_import import ModelImport + 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_import_config, + build_import_status_config, + ) + + +class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API, ModelImport): + """ + Main activities orchestrator for the Laborious system. + + This class combines functionality from multiple activity classes to provide + a unified interface for all workflow operations. It manages database connections, + MLFlow model interactions, data quality validation, and OPC server communications. + + The class implements multiple inheritance to combine specialized functionality: + - Storage: Database operations and data persistence + - MLFlow: Model inference and transformation operations + - Gates: Data quality validation and filtering mechanisms + - OPC: Real-time data export to OPC servers + - ModelMetrics: Model performance metrics and drift detection + - API: PI Web API export operations for industrial systems + + Attributes: + postgres_config (dict): PostgreSQL connection configuration + mlflow_config (dict): MLFlow server configuration + opc_config (dict): OPC server configuration + pi_web_api_config (dict): PI Web API server configuration + logger (Logger): Logging and observability instance + notification_handler (NotificationHandler): Notification management instance + """ + + def __init__( + self, + postgres_config: dict[str, Any], + mlflow_config: dict[str, Any], + minio_config: dict[str, Any], + opc_config: dict[str, Any], + pi_web_api_config: dict[str, Any], + logger: Logger, + notification_handler: NotificationHandler, + mongo_config: dict[str, Any] | None = None, + import_config: dict[str, Any] | None = None, + import_status_config: dict[str, Any] | None = None, + ): + """ + Initialize the Activities orchestrator with all required configurations. + + This constructor initializes all parent classes with their respective + configurations and sets up the foundation for all activity operations. + + Args: + postgres_config: PostgreSQL connection configuration dictionary + Required keys: host, port, user, password, dbname, min_connections, max_connections + mlflow_config: MLFlow server configuration dictionary + Required keys: host, port, username, password + opc_config: OPC server configuration dictionary + Can contain multiple server configurations + pi_web_api_config: PI Web API server configuration dictionary + Required keys: base_url, auth_type, auth_token + logger: Logger instance for observability and debugging + notification_handler: Notification handler for alerts and monitoring + mongo_config: MongoDB configuration for the model listing the import writes. When + absent, the import activities that need it refuse rather than guess — the + notification handler keeps its own client either way. + import_config: `build_import_config()`; read from the environment when absent + import_status_config: `build_import_status_config()`; read from the environment when + absent. This is the BFF database holding the import log, a second Postgres + connection that is never the one Storage holds. + + Raises: + Exception: If any parent class initialization fails + """ + metrics_controller = MetricsController(logger=logger) + + minio_repository = MinioRepository( + endpoint=minio_config['endpoint_url'], + access_key=minio_config['access_key'], + secret_key=minio_config['secret_key'], + bucket=minio_config['default_bucket'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + secure=minio_config['secure'], + ) + + # Initialize parent classes + Storage.__init__( + self, + host=postgres_config['host'], + port=postgres_config['port'], + user=postgres_config['user'], + password=postgres_config['password'], + dbname=postgres_config['dbname'], + min_connections=postgres_config['min_connections'], + max_connections=postgres_config['max_connections'], + retention_hours=minio_config['retention_hours'], + minio_repository=minio_repository, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + MLFlow.__init__( + self, + mlflow_host=mlflow_config['host'], + mlflow_port=mlflow_config['port'], + mlflow_username=mlflow_config['username'], + mlflow_password=mlflow_config['password'], + minio_repository=minio_repository, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + Gates.__init__( + self, + minio_repository=minio_repository, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + OPC.__init__( + self, + opc_servers=opc_config, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + ModelMetrics.__init__( + self, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + API.__init__( + self, + base_url=pi_web_api_config['base_url'], + auth_type=pi_web_api_config['auth_type'], + auth_token=pi_web_api_config['auth_token'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + mongo_repository = ( + MongoDBRepository( + connection_string=mongo_config['connection_string'], + database_name=mongo_config['database_name'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + if mongo_config + else None + ) + + ModelImport.__init__( + self, + status_config=import_status_config or build_import_status_config(), + import_config=import_config or build_import_config(), + minio_repository=minio_repository, + # The MLflow repository is shared rather than rebuilt: setting the tracking URI is a + # process-wide side effect, so two instances would be two chances to disagree. + mlflow_repository=self.model_monitoring_repository, + mongo_repository=mongo_repository, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + async def shutdown(self): + """ + Gracefully shutdown all activities and clean up resources. + + This method ensures proper cleanup of all resources including: + - PostgreSQL connection pools + - OPC server connections + - PI Web API client connections + - MLFlow model repositories + - Any other resources that need explicit cleanup + + The method should be called before the application terminates to ensure + proper resource cleanup and prevent resource leaks. + """ + Storage.close(self) + MLFlow.close(self) + Gates.close(self) + await OPC.aclose(self) + ModelMetrics.close(self) + API.close(self) + ModelImport.close(self) diff --git a/laborious/activities/api.py b/laborious/activities/api.py new file mode 100644 index 0000000..0e4114f --- /dev/null +++ b/laborious/activities/api.py @@ -0,0 +1,305 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import json + import traceback + from typing import Any + + from pandas import DataFrame + from sientia_do.notifications.handlers import NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.pi_web_api_client import PIWebAPIClient + + from laborious import metrics + + +PI_WEB_API_PREDICTION_ERROR_CONFIDENCE = 13 + + +class API(SientiaMonitoring): + """ + PI Web API operations for writing prediction data to PI Web API. + + This class provides Temporal activities for interacting with the PI Web API + to write prediction and confidence values to industrial systems. It handles + error scenarios gracefully by setting error confidence values and sending + notifications when write operations fail. + + The class implements comprehensive error handling for both prediction and + confidence value writes, ensuring that partial failures are properly + reported and handled. + """ + + def __init__( + self, + base_url: str, + auth_type: str, + auth_token: str, + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + ) -> None: + """ + Initialize API activity with PI Web API client. + + Args: + base_url (str): Base URL of the PI Web API server + auth_type (str): Authentication type ('basic' or 'bearer') + auth_token (str): Authentication token + logger (Logger): Logger instance for operation logging + notification_handler (NotificationHandler): Handler for system notifications + metrics_controller (MetricsController): Controller for metrics collection + """ + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + self.pi_web_api_client = PIWebAPIClient( + base_url=base_url, + auth_config={ + 'type': auth_type, + 'token': auth_token, + }, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + headers_config={ + 'Content-Type': 'application/json', + 'Accept': 'application/json', + 'x-requested-with': 'piwebapistreams', + 'User-Agent': 'Aig-Laborious-Agent/1.0', + }, + ) + + def get_pi_web_api_core_labels( + self, + metadata: dict[str, Any], + operation_type: str = 'write_pi_web_api_data', + ) -> dict[str, Any]: + """ + Generate core labels for PI Web API metrics. + + PI Web API metrics in laborious use the shared ``CORE_LABELS`` from + ``sientia_do``, which includes ``operation_type``. For this reason, + operation_type must always be present in emitted labels. + + Args: + - metadata (dict[str, Any]): Workflow execution metadata used to derive labels. + - operation_type (str): Operation type label for metric cardinality. + + Return: + dict[str, Any]: Core labels dictionary including operation_type. + """ + return super().get_core_labels( + metadata=metadata, + operation_type=operation_type, + ) + + def close(self) -> None: + """ + Close the PI Web API client and shutdown monitoring services. + + This method properly closes all connections and resources associated + with the PI Web API client and monitoring services. + """ + self.pi_web_api_client.close() + SientiaMonitoring.shutdown(self) + + async def process_pi_web_api_response( + self, + response_data: list[dict[str, Any]], + tags: dict[str, str], + core_labels: dict[str, str], + metadata: dict[str, Any], + ) -> tuple[int, str]: + """ + Process the response data from PI Web API write operation. + + Validates that all tags were successfully written, emits metrics for each tag + (success or error), and returns the appropriate prediction confidence value. + Sets error confidence if any tag write fails or if the number of written tags + doesn't match the expected count. + + Args: + - response_data (dict[str, Any]): The response data from the PI Web API write operation. + - tags (dict[str, str]): The tags that were written to the PI Web API. + - core_labels (dict[str, str]): The core labels of the workflow execution. + - metadata (dict[str, Any]): The metadata of the workflow execution. + Returns: + int: Prediction confidence value (0 for success, 13 for errors) + """ + + # Convert tags from name:webid to webid:name + tags = {w: t for t, w in tags.items()} + + tag_names = list[str](tags.values()) + + confidence = 0 + + message = '' + + # Evaluate response for each tag + written_tags = [] + for item in response_data: + web_id = item.get('WebId') + if not web_id: + self.error('The response did not contain some WebIds', metadata) + continue + errors = item.get('Errors', []) + tag_name = tags.get(web_id) + if not tag_name: + self.error( + f'The response did not contain the tag name for WebId {web_id}', metadata + ) + continue + if errors: + self.error( + f'Error writing tag {tag_name}:{web_id} to PI Web API: {errors}', metadata + ) + await self.emit_metric( + metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT, + tags={ + **core_labels, + 'tag_name': tag_name, + }, + ) + confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE + else: + await self.emit_metric( + metric_object=metrics.PI_WEB_API_PREDICTION_WRITTEN_COUNT, + tags={ + **core_labels, + 'tag_name': tag_name, + }, + ) + written_tags.append(tag_name) + + if len(written_tags) != len(tag_names): + message = f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.' + + self.error( + f'{message}\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}', + metadata, + ) + + await self.send_notification_async( + metadata=metadata, + notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR', + message=f'The number of written tags does not match the number of tag names: Expected {tag_names} tags, but {written_tags} tags were written.\nResponse:\n {json.dumps(response_data, indent=4)}\nTags:\n {json.dumps(tags, indent=4)}', + block='write_pi_web_api_data', + level=NotificationLevel.ERROR, + ) + confidence = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE + + return confidence, message + + @activity.defn(name='write_pi_web_api_data') + async def write_pi_web_api_data(self, input_data: dict[str, Any]) -> dict[Any, Any]: + """ + Write prediction and confidence data to PI Web API. + + Writes prediction values and confidence scores to PI Web API using configured + web IDs. Processes responses to validate writes and emit metrics. Handles errors + gracefully by setting error confidence values when writes fail and sending + notifications for both prediction and confidence write errors. + + Args: + input_data (dict[str, Any]): The input data containing: + - metadata (dict[str, Any]): Workflow execution metadata + - pi_web_api_output_config (dict[str, Any]): PI Web API configuration with: + - endpoint (str): PI Web API endpoint URL + - prediction_tags (dict[str, str]): Mapping of tag names to web IDs for predictions + - confidence_tags (dict[str, str]): Mapping of tag names to web IDs for confidence + - data (dict[str, Any]): Prediction data, its a dataframe converted to dict. + Returns: + dict[Any, Any]: Data dictionary with potentially modified confidence values + If prediction write fails, prediction_confidence is set to error value (13) + """ + metadata = input_data['metadata'] + data = DataFrame(input_data['data']) + pi_web_api_output_config = input_data['pi_web_api_output_config'] + + self.info(f'Writing data to PI Web API... config: {pi_web_api_output_config}', metadata) + + raw_prediction_tags = pi_web_api_output_config['prediction_tags'] + raw_confidence_tags = pi_web_api_output_config['confidence_tags'] + prediction_tags = list[str](raw_prediction_tags.values()) + confidence_tags = list(raw_confidence_tags.values()) + + core_labels = self.get_pi_web_api_core_labels(metadata) + + prediction_value = data.head(1)['prediction'].values[0] + confidence_value = data.head(1)['prediction_confidence'].values[0] + + try: + prediction_response = await self.pi_web_api_client.write_value( + web_ids=prediction_tags, + value={ + 'Timestamp': data.head(1)['timestamp'].values[0], + 'Value': prediction_value, + }, + metadata=metadata, + ) + + confidence, message = await self.process_pi_web_api_response( + response_data=prediction_response, + tags=raw_prediction_tags, + core_labels=core_labels, + metadata=metadata, + ) + + # Preserve incoming confidence/comments on successful PI writes. + # Only downgrade confidence or override comments when PI response + # explicitly reports a problem (e.g. partial write mismatch). + if confidence != 0: + data['prediction_confidence'] = confidence + if message: + data['comments'] = message + + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR', + message=f'Error writing prediction data to PI Web API: {e}\n Tags: {raw_prediction_tags}', + block='write_pi_web_api_data', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + self.error(trace, metadata) + + data['prediction_confidence'] = PI_WEB_API_PREDICTION_ERROR_CONFIDENCE + data['comments'] = str(e) + + return data.to_dict() + + try: + confidence_response = await self.pi_web_api_client.write_value( + web_ids=confidence_tags, + value={ + 'Timestamp': data.head(1)['timestamp'].values[0], + 'Value': float(confidence_value), + }, + metadata=metadata, + ) + + await self.process_pi_web_api_response( + response_data=confidence_response, + tags=raw_confidence_tags, + core_labels=core_labels, + metadata=metadata, + ) + + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR', + message=f'Error writing confidence data to PI Web API: {e}\n Tags: {raw_confidence_tags}', + block='write_pi_web_api_data', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + return data.to_dict() diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py new file mode 100644 index 0000000..ad88ca1 --- /dev/null +++ b/laborious/activities/gates.py @@ -0,0 +1,804 @@ +from sientia_do.repository.minio_repository import MinioRepository +from temporalio import activity, workflow + +from laborious.utils.repository.minio_manager import MinioManager + +with workflow.unsafe.imports_passed_through(): + import traceback + from collections.abc import Callable, Mapping + from typing import Any + + from pandas import DataFrame + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.utils.formatters import create_sample_dict + + from laborious import metrics + from laborious.utils.dataframe_debug import build_dataframe_debug_message + from laborious.utils.filters.conditional_filters import ( + filter_empty_data, + filter_specific_variables_null_values, + ) + from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + +# Strongly-typed filter function signatures +InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool] +ResponseFilterFunc = Callable[[dict[str, Any], dict[str, Any]], bool] +ContentFilterFunc = Callable[[DataFrame, dict[str, Any]], bool] + +# Input filter function mappings +input_filter_functions: dict[str, InputFilterFunc] = { + 'SPECIFIC_VARIABLES_NULL_VALUES': filter_specific_variables_null_values, + 'EMPTY_DATA': filter_empty_data, +} + +# Confidence mappings kept separate from function maps to avoid Union types +input_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 2, + 'REPEAT': -1, +} + +# MLFlow response filter function mappings +mlflow_response_filter_functions: dict[str, ResponseFilterFunc] = { + 'API_ERROR': api_error_filter, +} + +mlflow_response_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 10, + 'REPEAT': -1, +} + +# MLFlow content filter function mappings +mlflow_content_filter_functions: dict[str, ContentFilterFunc] = { + 'NAN_VALUES': nan_values_filter, + 'EMPTY_DATA': filter_empty_data, +} + +mlflow_content_path_confidence: Mapping[str, int] = { + 'STOP': -1, + 'CONTINUE': 18, + 'REPEAT': -1, +} + + +class Gates(MinioManager): + """ + Data quality gates and filtering activities for the Laborious system. + + This class implements comprehensive data quality validation and filtering + mechanisms that can be applied at different stages of the prediction pipeline. + It provides configurable filters with policy-based decision making to ensure + data integrity and quality throughout the ML workflow. + + The class supports multiple filter types and implements a flexible policy + system that can be configured for different validation requirements. Each + filter returns a path decision (STOP, CONTINUE, REPEAT) along with confidence + scores and detailed comments for monitoring and debugging. + + Attributes: + input_filter_functions (dict): Mapping of input filter names to functions + mlflow_response_filter_functions (dict): Mapping of MLFlow response filter names to functions + mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions + """ + + minio_repository: MinioRepository | None = None + _MAX_DEBUG_DATAFRAME_ROWS = 100 + + def __init__( + self, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + """ + Initialize data quality gates with logging and notification capabilities. + + Args: + logger: Logger instance for observability and debugging + notification_handler: Notification handler for alerts and monitoring + + Raises: + Exception: If BaseActivity initialization fails + """ + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) + + def close(self) -> None: + """ + Close the gates activity and clean up resources. + """ + + MinioManager.close(self) + + def __del__(self): + self.close() + + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like payload to be logged + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + + @staticmethod + def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]: + """ + Read filter policy/config keys in a case-insensitive way. + + Args: + config (dict[str, Any]): Filter configuration dictionary. + + Return: + tuple[str, dict[str, Any]]: Parsed policy and config payload. + """ + normalized = {str(key).upper(): value for key, value in config.items()} + policy = normalized['POLICY'] + filter_config = normalized.get('CONFIG', {}) + return policy, filter_config + + @activity.defn(name='input_gate') + async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: + """ + Apply input data quality filters and validation. + + This activity validates input data quality using configurable filters + before proceeding with ML operations. It applies multiple filter types + and returns a path decision based on the filter results and configured + policies. + + The method implements a comprehensive filtering system that: + 1. Applies configured filters to input data + 2. Evaluates filter results against policy configurations + 3. Determines appropriate path decisions (STOP, CONTINUE, REPEAT) + 4. Provides confidence scores and detailed comments + 5. Handles errors gracefully with notification integration + + Args: + input_data: Configuration and data for input validation + Required keys: + - metadata (dict): Workflow execution metadata + - filters (dict): Filter configuration and policies + - data (dict): Input data to validate + - path_priority (list[str]): Priority order for path decisions + + Returns: + tuple: (path_flag, confidence, comment) + - path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None) + - confidence (int): Confidence score for the decision + - comment (str): Detailed explanation of the decision + + Raises: + Exception: If filter execution fails or configuration is invalid + """ + metadata = input_data['metadata'] + + self.info('Performing input gate...', metadata) + + filters = input_data['filters'] + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + path_priority = input_data['path_priority'] + + filter_output = [] + + self._debug_dataframe('Input data:', data, metadata) + self.debug(f'Filters: {filters}', metadata) + + # Apply each configured filter + for fil, config in filters.items(): + if fil not in input_filter_functions: + self.error(f'Filter {fil} not found', metadata) + continue + policy, filter_config = self._read_filter_entry(config) + try: + if input_filter_functions[fil](data, filter_config): + self.debug(f'Data not passed the input filter {fil}:{config}', metadata) + filter_output.append(policy) + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id=f'INTPUT_GATE_ERROR__{fil}', + message=f'Error in filter {fil}:{config}: \n {e}', + block='input_gate', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + for path_flag in path_priority: + if path_flag in filter_output: + self.info(f'Input gate result: {path_flag}', metadata) + return path_flag, input_path_confidence[path_flag], 'Input data with bad quality' + + self.info('Nothing was filtered by the input gate', metadata) + + del data + + return None, 0, '' + + @activity.defn(name='mlflow_response_gate') + async def mlflow_response_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: + """ + Validate MLFlow API response quality and integrity. + + This activity validates MLFlow API responses to ensure they meet quality + standards before proceeding with further processing. It applies response-specific + filters and determines appropriate path decisions based on response quality. + + The method implements response validation that: + 1. Applies MLFlow response-specific filters + 2. Evaluates API response quality and integrity + 3. Determines path decisions based on response validation results + 4. Provides confidence scores and detailed validation comments + 5. Handles API errors and response validation failures + + Args: + input_data: Configuration and data for response validation + Required keys: + - metadata (dict): Workflow execution metadata + - filters (dict): Response filter configuration and policies + - data (dict): MLFlow API response data to validate + - type (str): Type of MLFlow operation (transform, predict) + - path_priority (list[str]): Priority order for path decisions + + Returns: + tuple: (path_flag, confidence, comment) + - path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None) + - confidence (int): Confidence score for the decision + - comment (str): Detailed explanation of the decision + + Raises: + Exception: If response validation fails or configuration is invalid + """ + metadata = input_data['metadata'] + self.info('Performing mlflow response gate...', metadata) + raw_data = input_data['data'] + filters = input_data['filters'] + + self.debug( + f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata + ) + self.debug(f'Filters: {filters}', metadata) + + payload = MinioDataFramePayload.from_dict(raw_data) + data = await payload.retrieve(self.minio_repository, metadata) + + gate_type = input_data['type'] + path_priority = input_data['path_priority'] + + filter_output = [] + + comments = [] + + status = payload.status or {} + + for fil, config in filters.items(): + if fil not in mlflow_response_filter_functions: + continue + policy, filter_config = self._read_filter_entry(config) + try: + if mlflow_response_filter_functions[fil](status, filter_config): + filter_output.append(policy) + comments.append(status.get('message', 'Unknown MLFlow API error')) + await self.send_notification_async( + metadata=metadata, + notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}', + message=status.get('message', 'Unknown MLFlow API error'), + block='mlflow_gate', + level=NotificationLevel.ERROR, + attachment_content=status.get('traceback'), + ) + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}', + message=f'Error in filter {fil}:{config}: \n {e}', + block='mlflow_gate', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + for path_flag in path_priority: + if path_flag in filter_output: + self.info(f'Mlflow response gate result: {path_flag}', metadata) + return path_flag, mlflow_response_path_confidence[path_flag], ', '.join(comments) + + self.info('Nothing was filtered by the mlflow response gate', metadata) + + del data + + return None, 0, '' + + @activity.defn(name='mlflow_content_gate') + async def mlflow_content_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]: + """ + Validate MLFlow prediction content quality and integrity. + + This activity validates the content of MLFlow predictions to ensure they + meet quality standards before export and persistence. It applies content-specific + filters and determines appropriate path decisions based on content quality. + + The method implements content validation that: + 1. Applies MLFlow content-specific filters + 2. Evaluates prediction content quality and integrity + 3. Determines path decisions based on content validation results + 4. Provides confidence scores and detailed validation comments + 5. Handles content validation failures and quality issues + + Args: + input_data: Configuration and data for content validation + Required keys: + - metadata (dict): Workflow execution metadata + - filters (dict): Content filter configuration and policies + - data (dict): MLFlow prediction content to validate + - type (str): Type of MLFlow operation (transform, predict) + - path_priority (list[str]): Priority order for path decisions + + Returns: + tuple: (path_flag, confidence, comment) + - path_flag (str | None): Decision path (STOP, CONTINUE, REPEAT, or None) + - confidence (int): Confidence score for the decision + - comment (str): Detailed explanation of the decision + + Raises: + Exception: If content validation fails or configuration is invalid + """ + metadata = input_data['metadata'] + self.info('Performing mlflow content gate...', metadata) + + filters = input_data['filters'] + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + gate_type = input_data['type'] + path_priority = input_data['path_priority'] + + filter_output = [] + + self._debug_dataframe('Input data:', data, metadata) + self.debug(f'Filters: \n {filters}', metadata) + + for fil, config in filters.items(): + if fil not in mlflow_content_filter_functions: + continue + policy, filter_config = self._read_filter_entry(config) + try: + if mlflow_content_filter_functions[fil](data, filter_config): + filter_output.append(policy) + await self.send_notification_async( + metadata=metadata, + notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}', + message=f'Data not passed the content filter {fil}:{config}', + block='mlflow_gate', + level=NotificationLevel.WARNING, + attachment_content=data.to_string(), + ) + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}', + message=f'Error in filter {fil}:{config}: \n {e}', + block='mlflow_gate', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + for path_flag in path_priority: + if path_flag in filter_output: + self.info(f'Mlflow content gate result: {path_flag}', metadata) + return ( + path_flag, + mlflow_content_path_confidence[path_flag], + 'Transformed data not passed the content filter', + ) + + self.info('Nothing was filtered by the mlflow content gate', metadata) + + del data + + return None, 0, '' + + def get_prediction_store_policy( + self, prediction_store_policy: str, metadata: dict[str, Any] + ) -> tuple[str, int]: + """ + Parse and validate prediction store policy configuration. + + This method parses prediction store policy strings in the format 'type:value' + and validates them against allowed policy types and values. It provides + sensible defaults for invalid configurations and logs policy validation + failures for operational monitoring. + + Supported Policy Types: + - 'lts': Latest timestamp - sorts data by timestamp descending + - 'erl': Earliest timestamp - sorts data by timestamp ascending + + Args: + prediction_store_policy (str): Policy string in format 'type:value' + metadata (dict[str, Any]): Context metadata for logging and notifications + + Returns: + tuple[str, int]: (policy_type, policy_value) + - policy_type (str): Validated policy type ('lts' or 'erl') + - policy_value (int): Number of rows to retain + """ + policy_elements = prediction_store_policy.split(':') + + if len(policy_elements) < 2: + self.error( + f'Invalid prediction store policy: {prediction_store_policy}, using default policy', + metadata, + ) + return 'lts', 1 + + policy_type = policy_elements[0] + policy_value = policy_elements[1] + + # If the policy_type is not lts or erl, we use the default policy + # If the policty_value is not a number or 0, we use the default policy + if ( + policy_type not in ['lts', 'erl'] + or not policy_value.isdigit() + or int(policy_value) == 0 + ): + self.error( + f'Invalid prediction store policy: {prediction_store_policy}, using default policy', + metadata, + ) + return 'lts', 1 + + return policy_type, int(policy_value) + + @activity.defn(name='format_transformed_data') + async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload: + """ + Format transformed data for storage and export operations. + + This method formats transformed data from MLFlow model transformations + into a standardized format suitable for database storage. It converts + wide-format data (columns as variables) into long-format (melted) + with proper timestamp handling and model identification. + + The formatting process includes: + 1. Converting input data dictionary to DataFrame + 2. Extracting timestamps from DataFrame index + 3. Resetting index to create sequential row numbers + 4. Melting data from wide format to long format (variable-value pairs) + 5. Adding model_id for data lineage tracking + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - data (dict[str, Any]): Transformed data to format (DataFrame-compatible dict) + - model_id (str): Unique identifier for the ML model + + Returns: + dict: Formatted data dictionary with keys: + - timestamp (dict): Timestamp values indexed by row number + - variable (dict): Variable names indexed by row number + - value (dict): Variable values indexed by row number + - model_id (dict): Model identifiers indexed by row number + """ + metadata = input_data['metadata'] + + model_id = input_data['model_id'] + + self.info('Formatting transformed data...', metadata) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + data['timestamp'] = data.index + data = data.reset_index(drop=True) + + data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') + data['model_id'] = model_id + + return await MinioDataFramePayload.from_dataframe( + dataframe=data, + minio_repo=self.minio_repository, + model_name=input_data['model_name'], + operation='transform', + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + @activity.defn(name='format_prediction') + async def format_prediction(self, input_data: dict[str, Any]) -> dict: + """ + Format prediction data according to configured storage policies. + + This method formats prediction data for storage and export operations. + It applies timestamp-based sorting policies, adds metadata fields, + and ensures data consistency before persistence. The method supports + multiple storage policies for flexible data retention strategies. + + If only one row is present, we use the last timestamp as the timestamp + + Storage Policies: + - 'lts:N': Latest timestamp - retains N most recent predictions + - 'erl:N': Earliest timestamp - retains N oldest predictions + + Args: + input_data (dict): Input data containing: + - data (dict[str, Any]): Raw prediction data to format + - timestamp (str): Timestamp of the data + - model_id (str): Unique identifier for the ML model + - prediction_confidence (float): Confidence score for the prediction + - prediction_store_policy (str): Storage policy in format 'type:value' + + Returns: + dict: Formatted prediction data ready for storage and export + """ + metadata = input_data['metadata'] + last_timestamp = input_data['timestamp'] + prediction_store_policy = input_data['prediction_store_policy'] + self.info('Formatting prediction...', metadata) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + # Create timestamp column from index and reset index + data['timestamp'] = data.index + data = data.reset_index(drop=True) + + self.debug(f'Prediction store policy: {prediction_store_policy}', metadata) + self._debug_dataframe('Prediction data:', data, metadata) + + policy_type, policy_value = self.get_prediction_store_policy( + prediction_store_policy, metadata + ) + + # If data has no timestamp, we use the default timestamp and not sort the data + self.info( + f'Sorting data by timestamp and applying policy: {policy_type}:{policy_value}', metadata + ) + + # If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows + if policy_type == 'lts': + self.debug('Sorting data by timestamp descending', metadata) + data = data.sort_values(by='timestamp', ascending=False) + # If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows + elif policy_type == 'erl': + self.debug('Sorting data by timestamp ascending', metadata) + data = data.sort_values(by='timestamp', ascending=True) + else: + self.error(f'Invalid policy type: {policy_type}, using default policy', metadata) + raise ValueError(f'Invalid policy type: {policy_type}') + + int_policy_value = int(policy_value) + + data = data.head(int_policy_value) + if int_policy_value == 1: + data['timestamp'] = last_timestamp + + data['model_id'] = input_data['model_id'] + data['prediction_confidence'] = input_data['prediction_confidence'] + data['prediction_status'] = 'Good' + data['comments'] = '' + data = data.sort_values(by='timestamp', ascending=False) + data = data.reset_index(drop=True) + + self.info(f'Prediction formatted: {len(data)} rows', metadata) + self._debug_dataframe('Prediction data:', data, metadata) + + return data.to_dict() + + @activity.defn(name='format_default_prediction') + async def format_default_prediction(self, input_data: dict[str, Any]) -> dict: + """ + Create and format default prediction data for error conditions. + + This method generates default prediction data when the main prediction + pipeline encounters errors or quality issues. It creates a standardized + data structure with zero values for predictions and useful metadata + for operational monitoring and debugging. + + The default prediction serves as a fallback mechanism to: + 1. Maintain data pipeline continuity during failures + 2. Provide operational visibility into prediction quality issues + 3. Enable downstream systems to handle error conditions gracefully + 4. Support debugging and troubleshooting efforts + + Args: + input_data (dict): Input data containing: + - timestamp (str): Timestamp for the default prediction + - model_id (str): Unique identifier for the ML model + - prediction_confidence (float): Confidence score (typically low for errors) + - comment (str): Error description or operational comment + + Returns: + dict: Formatted default prediction data with error indicators + """ + + metadata = input_data['metadata'] + self.debug('Formatting default prediction...', metadata) + + data = DataFrame( + { + 'prediction': [0], + 'response_time': [0], + 'timestamp': [input_data['timestamp']], + 'model_id': [input_data['model_id']], + 'prediction_confidence': [input_data['prediction_confidence']], + 'prediction_status': ['Bad'], + 'comments': [input_data['comment']], + } + ) + + self.info(f'Default prediction formatted: {data.size} rows', metadata) + return data.to_dict() + + @activity.defn(name='format_retrain_report') + async def format_retrain_report(self, input_data: dict[str, Any]) -> dict: + """ + Format retrain report data for storage and audit trail maintenance. + + This method formats model retraining operation results into a standardized + report format suitable for database storage and operational monitoring. + It captures retraining status, timestamps, and model version information + for comprehensive audit trails and operational visibility. + + The formatting process includes: + 1. Extracting retraining experiment response data + 2. Capturing model update report information (version, MLflow IDs) + 3. Formatting timestamps and status information + 4. Conditionally including version information for successful retrains + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - experiment_response (dict): Retraining experiment response containing: + - success (bool): Retraining operation success status + - timestamp (str): Timestamp of the retraining operation + - message (str): Status message or error description + - update_report (dict): Model update report containing: + - version (str): New model version identifier + - mlflow_run_id (str): MLflow run identifier + - mlflow_experiment_id (str): MLflow experiment identifier + - model_id (str): Unique identifier for the ML model + - model_name (str): Name of the ML model + + Returns: + dict: Formatted retrain report dictionary with keys: + - model_id (dict): Model identifiers indexed by row number + - model_name (dict): Model names indexed by row number + - timestamp (dict): Retraining timestamps indexed by row number + - status (dict): Retraining status messages indexed by row number + - version (dict, optional): Model versions indexed by row number + Only included if experiment_response['success'] is True + - mlflow_run_id (dict, optional): MLflow run IDs indexed by row number + Only included if experiment_response['success'] is True + - mlflow_experiment_id (dict, optional): MLflow experiment IDs indexed by row number + Only included if experiment_response['success'] is True + """ + metadata = input_data['metadata'] + self.info('Formatting retrain report...', metadata) + + experiment_response = input_data['experiment_response'] + update_report = input_data['update_report'] + model_id = input_data['model_id'] + model_name = input_data['model_name'] + + report = DataFrame( + { + 'model_id': [model_id], + 'model_name': [model_name], + 'timestamp': [experiment_response['timestamp']], + 'status': [experiment_response['message']], + } + ) + + if experiment_response['success']: + # Retrain was successfull + report['version'] = update_report['version'] + report['mlflow_run_id'] = update_report['mlflow_run_id'] + report['mlflow_experiment_id'] = update_report['mlflow_experiment_id'] + + self._debug_dataframe('Retrain report:', report, metadata) + + return report.to_dict() + + @activity.defn(name='write_metrics') + async def write_metrics(self, input_data: dict[str, Any]): + """ + Write prediction performance metrics to Prometheus monitoring system. + + This method records comprehensive metrics for prediction operations, + enabling operational monitoring, performance analysis, and alerting. + It tracks prediction counts, confidence levels, and response times + for each model and pipeline combination. + + Metrics Recorded: + 1. Prediction Count: Incremental counter for successful predictions + 2. Confidence Monitor: Current confidence level for predictions + 3. Response Time Monitor: Histogram of prediction response times + + Args: + input_data (dict): Input data containing: + - metadata (dict[str, Any]): Workflow execution metadata + - prediction (dict[str, Any]): Prediction data with metrics + + Raises: + Exception: If metrics writing fails or configuration is invalid + """ + metadata = input_data['metadata'] + prediction = DataFrame(input_data['prediction']) + prediction_confidence = prediction['prediction_confidence'].values[0] + response_time = prediction['response_time'].values[0] + opc_metrics = input_data['opc_metrics'] + + self.info(f'Writing metrics for model {metadata["model_name"]}', metadata) + + core_tags = { + 'pod_id': self.pod_id, + 'runtime': self.runtime, + 'operation_type': 'predict', + 'model_name': metadata['model_name'], + 'workflow_name': metadata['workflow_name'], + } + await self.emit_metric( + metric_object=metrics.PREDICTIONS_WRITTEN_COUNT, + tags=core_tags, + ) + + await self.emit_metric( + metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR, + method='set', + tags=core_tags, + value=prediction_confidence, + ) + + await self.emit_metric( + metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR, + method='observe', + tags=core_tags, + value=response_time, + ) + + for server_id, tags in opc_metrics.items(): + for tag, response_time in tags.items(): + if response_time is not None: + await self.emit_metric( + metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, + method='observe', + tags={ + **core_tags, + 'opc_server_id': server_id, + 'tag': tag, + }, + value=response_time, + ) + + await self.emit_metric( + metric_object=metrics.PREDICTION_OPC_WRITING_COUNT, + tags={ + **core_tags, + 'opc_server_id': server_id, + 'tag': tag, + }, + ) + + self.info(f'Metrics written for model {metadata["model_name"]}', metadata) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py new file mode 100644 index 0000000..f4d42a3 --- /dev/null +++ b/laborious/activities/mlflow.py @@ -0,0 +1,528 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import traceback + from typing import Any + + import numpy as np + from pandas import to_datetime + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.temporal.constants import ( + DATETIME_FORMAT, + DATETIME_FORMAT_MS_WITH_TZ, + DATETIME_FORMAT_WITH_TZ, + now, + ) + from sientia_do.utils.formatters import create_sample_dict + + from laborious.utils.dataframe_debug import build_dataframe_debug_message + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + from laborious.utils.repository.minio_manager import MinioManager + from laborious.utils.repository.model_repository import MLFlowRepository + + +class MLFlow(MinioManager): + """ + MLFlow integration activities for model inference operations. + + This class provides activities for interacting with MLFlow models, including + data transformation and prediction operations. It handles authentication, + data preprocessing, and model management with configurable retention policies. + + The class implements comprehensive error handling and logging for all + MLFlow operations, ensuring reliable model inference in production environments. + + Attributes: + mlflow_host (str): MLFlow server hostname + mlflow_port (int): MLFlow server port + mlflow_username (str): MLFlow authentication username + mlflow_password (str): MLFlow authentication password + model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations + """ + + _MAX_DEBUG_DATAFRAME_ROWS = 100 + + def __init__( + self, + mlflow_host: str, + mlflow_port: int, + mlflow_username: str, + mlflow_password: str, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + """ + Initialize MLFlow activities with server configuration. + + Args: + mlflow_host: MLFlow server hostname or IP address + mlflow_port: MLFlow server port number + mlflow_username: Username for MLFlow authentication + mlflow_password: Password for MLFlow authentication + logger: Logger instance for observability and debugging + notification_handler: Notification handler for alerts and monitoring + + Raises: + Exception: If MLFlowRepository initialization fails + """ + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) + self.mlflow_host = mlflow_host + self.mlflow_port = mlflow_port + self.mlflow_username = mlflow_username + self.mlflow_password = mlflow_password + + self.model_monitoring_repository = MLFlowRepository( + f'{mlflow_host}:{mlflow_port}', + mlflow_username, + mlflow_password, + logger, + notification_handler, + metrics_controller, + ) + + def close(self) -> None: + """ + Close the MLFlow activity and clean up resources. + """ + MinioManager.close(self) + + def __del__(self): + self.close() + + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like object expected to expose shape and to_csv + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + + @activity.defn(name='request_transform') + async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload: + """ + Transform input data using MLFlow models. + + This activity processes input data through MLFlow model transformation, + including data preprocessing, format conversion, and validation. It handles + data deduplication, pivoting, and cleanup to ensure optimal model performance. + + The transformation process includes: + 1. Data deduplication based on variable and timestamp + 2. Data pivoting for model input format + 3. Null value handling and cleanup + 4. MLFlow model transformation request + 5. Response validation and logging + + Args: + input_data: Configuration and data for transformation + Required keys: + - metadata (dict): Workflow execution metadata + - data (dict): Input data for transformation + - model_name (str): Name of the MLFlow model to use + - model_retention (int): Model retention period in minutes + + Returns: + dict: Transformed data from MLFlow model + + Raises: + Exception: If transformation fails or MLFlow model is unavailable + """ + metadata = input_data['metadata'] + self.info('Transforming data...', metadata) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + + self._debug_dataframe('Raw input data:', data, metadata) + + # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair + data = data.sort_values('created_at', ascending=False).drop_duplicates( + subset=['variable', 'timestamp'], keep='first' + ) + + # Pivot data for model input format + data = data.pivot(index='timestamp', columns='variable', values='value') + data.fillna(np.nan, inplace=True) + + data.columns.name = None + data.index.name = None + + data['timestamp'] = data.index + + self._debug_dataframe('Processed input data:', data, metadata) + + # Request transformation from MLFlow model + response_data = await self.model_monitoring_repository.transform( + model_name, data, model_config, metadata + ) + + self.debug( + f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}', + metadata, + ) + + self.debug( + f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}', + metadata, + ) + + self.info('Data transformed successfully', metadata) + + if not response_data.get('success', False): + return await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=self.minio_repository, + model_name=model_name, + operation='transform', + status=response_data, + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + return await MinioDataFramePayload.from_dataframe( + dataframe=response_data['content'], + minio_repo=self.minio_repository, + model_name=model_name, + operation='transform', + workflow_metadata=metadata, + status={ + 'success': True, + }, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + @activity.defn(name='request_predict') + async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload: + """ + Execute predictions using MLFlow models. + + This activity performs ML model inference using MLFlow models with the + transformed data. It handles data format conversion, null value processing, + and model prediction requests with comprehensive error handling. + + The prediction process includes: + 1. Data format validation and cleanup + 2. Null value handling for model compatibility + 3. MLFlow model prediction request + 4. Response validation and logging + 5. Performance monitoring and metrics + + Args: + input_data: Configuration and data for prediction + Required keys: + - metadata (dict): Workflow execution metadata + - data (dict): Transformed data for prediction + - model_name (str): Name of the MLFlow model to use + - model_retention (int): Model retention period in minutes + + Returns: + dict: Prediction results from MLFlow model + + Raises: + Exception: If prediction fails or MLFlow model is unavailable + """ + metadata = input_data['metadata'] + self.info('Predicting data...', metadata) + + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + + self._debug_dataframe('Input data for prediction:', data, metadata) + + # Convert numpy.nan to None for model compatibility + data.replace(np.nan, None, inplace=True) + + data['timestamp'] = data.index + data['timestamp'] = to_datetime( + data['timestamp'], format=DATETIME_FORMAT_WITH_TZ + ).dt.strftime(DATETIME_FORMAT) + + # Request prediction from MLFlow model + response_data = await self.model_monitoring_repository.predict( + model_name, data, model_config, metadata + ) + + self.debug( + f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}', + metadata, + ) + + self.info('Data predicted successfully', metadata) + + if not response_data.get('success', False): + return await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=self.minio_repository, + model_name=model_name, + operation='predict', + status=response_data, + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + return await MinioDataFramePayload.from_dataframe( + dataframe=response_data['content'], + minio_repo=self.minio_repository, + model_name=model_name, + operation='predict', + workflow_metadata=metadata, + status={ + 'success': True, + }, + last_timestamp=payload.last_timestamp, + logger=self.logger, + ) + + @activity.defn(name='retrain_model') + async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Retrain MLFlow models with updated training data. + + This activity orchestrates the complete model retraining process, + including data preparation, model retraining execution, and result + validation. It handles data preprocessing, column cleanup, and + comprehensive error handling for production model management. + + The retraining process includes: + 1. Data timestamp extraction and validation + 2. Column cleanup and data preparation + 3. Data pivoting for model input format + 4. MLFlow model retraining execution + 5. Result validation and error handling + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - data (dict[str, Any]): Training data for model retraining + - model_name (str): Name of the MLFlow model to retrain + + Returns: + dict: Retraining results containing: + - status (str): Retraining operation status + - timestamp (str): Timestamp of the retraining operation + - experiment (str): MLFlow experiment identifier + + Raises: + Exception: If retraining fails or encounters critical errors + """ + + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + + metadata = input_data['metadata'] + + try: + # Payload-based retrain input (inline dict or MinIO offloaded). + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id='ERROR_LOADING_RETRAIN_DATA', + message=f'Error loading retrain data: {e}', + block='retrain_model', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata) + return { + 'success': False, + 'message': f'Error loading retrain data: {e}', + 'traceback': trace, + 'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ), + } + + self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata) + + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + + self.info(f'Retraining model {model_name}...', metadata) + + timestamp = data['timestamp'].max() + self.debug(f'Timestamp: {timestamp}', metadata) + + # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair + if 'created_at' in data.columns: + data = data.sort_values('created_at', ascending=False).drop_duplicates( + subset=['variable', 'timestamp'], keep='first' + ) + else: + data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first') + + data.drop(columns=['model_id'], inplace=True, errors='ignore') + data.drop(columns=['created_at'], inplace=True, errors='ignore') + + # Pivot data for model input format + data = data.pivot(index='timestamp', columns='variable', values='value') + data.fillna(np.nan, inplace=True) + # data.reset_index(inplace=True) + data.columns.name = None + + data['timestamp'] = data.index + data['timestamp'] = to_datetime( + data['timestamp'], format=DATETIME_FORMAT_WITH_TZ + ).dt.strftime(DATETIME_FORMAT) + data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT) + + data.columns.name = None + + retrain_output = await self.model_monitoring_repository.retrain_model( + data=data, model_name=model_name, model_config=model_config, metadata=metadata + ) + + if not retrain_output['success']: + trace = retrain_output['traceback'] + await self.send_notification_async( + metadata=metadata, + notification_id='RETRAIN_MODEL_ERROR', + message=f'Error retraining model {model_name}: {retrain_output["message"]}', + block='retrain_model', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata=metadata) + + return {**retrain_output, 'timestamp': timestamp} + + @activity.defn(name='update_production_model') + async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]: + """ + Update production model with newly trained model version. + + This activity manages the critical process of updating production + models with newly trained versions. It handles model deployment, + status tracking, and comprehensive reporting for operational + visibility and audit trails. + + The update process includes: + 1. Production model update execution + 2. Status and metadata tracking + 3. Comprehensive reporting and logging + 4. Error handling and notification + 5. Audit trail maintenance + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - model_name (str): Name of the MLFlow model to update + - experiment (str): MLFlow experiment identifier + - model_id (str): Unique identifier for the model version + - timestamp (str): Timestamp of the update operation + - status (str): Current status of the model update + + Returns: + dict[Any, Any]: Comprehensive update report containing: + - model_id (str): Model version identifier + - model_name (str): Name of the updated model + - timestamp (str): Update operation timestamp + - status (str): Update operation status + - Additional MLFlow response metadata + + Raises: + Exception: If production model update fails + """ + metadata = input_data['metadata'] + model_name = input_data['model_name'] + experiment = input_data['experiment'] + self.info( + f'Updating production model {model_name} from experiment {experiment}...', metadata + ) + + try: + response = await self.model_monitoring_repository.update_production_model( + experiment=experiment, model_name=model_name, metadata=metadata + ) + + self.info(f'Production model {model_name} updated successfully', metadata) + return response + + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id='UPDATE_PRODUCTION_MODEL_ERROR', + message=f'Error updating production model {model_name}: {e}', + block='update_production_model', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata=metadata) + raise e + + @activity.defn(name='get_reference_data') + async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None: + """ + Get reference data from the MLflow Model Registry. + + This method retrieves evaluation reference data stored as artifacts in the + MLflow Model Registry. The reference data is typically used for model + drift detection, performance comparison, and quality validation. The method + loads the data from a CSV artifact file and formats timestamps for + consistent processing. + + The method handles: + 1. Loading evaluation data artifact from MLflow Model Registry + 2. Timestamp parsing and formatting for consistency + 3. Data conversion to dictionary format for workflow consumption + 4. Graceful handling of missing reference data + + Args: + input_data (dict): Input data containing: + - metadata (dict): Workflow execution metadata + - model_name (str): Name of the MLFlow model to get reference data from + + Returns: + list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry + as a list of dictionaries. Returns None if reference data is not found + or if the artifact does not exist. + + Raises: + Exception: If artifact loading fails or encounters errors during processing + """ + + metadata = input_data['metadata'] + model_name = input_data['model_name'] + artifact = 'evaluation_data.csv' + + reference_data = await self.model_monitoring_repository.load_artifact_dataframe( + model_name=model_name, artifact_path=artifact, metadata=metadata + ) + + if reference_data is None: + self.warning(f'Reference data not found for model {model_name}', metadata) + return None + + reference_data['timestamp'] = to_datetime(reference_data['timestamp']) + reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT) + + return reference_data.to_dict(orient='records') diff --git a/laborious/activities/model_import.py b/laborious/activities/model_import.py new file mode 100644 index 0000000..125ee97 --- /dev/null +++ b/laborious/activities/model_import.py @@ -0,0 +1,1044 @@ +"""Model import activities: the import log row, the bundle, MLflow, the model document, cleanup. + +The import log row is **not created here**. It is `public.experiment_run` in the Postgres database +`sientia-core-mlops-bff`; the frontend inserts it through the BFF with `run_type='IMPORT'` and +`status='PENDING'`, and passes its `id` in the workflow input. Every write in this module is an +`UPDATE ... WHERE id = :import_run_id`, on an engine of its own against that database — never the +`sientia` database Laborious itself uses, and never in a transaction shared with a provisioning +write. + +Three classes of status write, and only the first one is strict: + +- **A, the claim** — `WHERE id AND run_type='IMPORT' AND (PENDING OR my own workflow id)`. This is + the ownership decision, and a claim matching nothing fails the import explicitly. +- **B, progress** — `WHERE id`, nothing else. A hint. Its failure never fails the import. +- **C, terminal** — `WHERE id`, nothing else, in the workflow's `finally`, after cleanup. The + authority: it rewrites status, code, reason and step from one classification of the outcome. +""" + +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import json + import shutil + from datetime import timedelta + from enum import StrEnum + from pathlib import Path + from typing import Any, Final + + from minio.commonconfig import ENABLED, Filter + from minio.lifecycleconfig import Expiration, LifecycleConfig, Rule + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + 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 import MinioRepository + from sientia_do.repository.mongodb_repository import MongoDBRepository + from sientia_do.temporal.constants import now + from sqlalchemy import create_engine, text + from sqlalchemy.engine import URL + from sqlalchemy.exc import DataError, IntegrityError, ProgrammingError + from temporalio.common import RetryPolicy + + from laborious import metrics + from laborious.activities.model_import_errors import ( + ImportErrorCode, + ImportGateRejectedError, + ImportLogRowNotClaimableError, + ImportStatusRejectedError, + classify_import_failure, + report_import_failure, + ) + from laborious.utils.bundle import format as fmt + from laborious.utils.bundle.reader import BundleGateError, BundleLimits, BundleReader + from laborious.utils.bundle.steps import ImportStep + from laborious.utils.import_password import ( + EnvelopeError, + decrypt_password_envelope, + load_envelope_key, + ) + from laborious.utils.repository.model_repository import MLFlowRepository + +# The table is a constant of the code, never an input: there is exactly one import log on this +# platform, and letting a caller name it would be a SQL-injection surface bought for nothing. The +# statements below write the name out literally rather than interpolating this constant — an +# interpolation is a shape a reader (and bandit) has to reason about, and there is nothing to +# parameterise. The constant is what the tests assert every statement against. +IMPORT_LOG_TABLE: Final[str] = 'public.experiment_run' + +RUN_NAME_MAX_LENGTH: Final[int] = 50 +RUN_NAME_PREFIX: Final[str] = 'import-' + +_MISSING_MINIO = 'Minio repository not initialized' +_MISSING_MLFLOW = 'MLflow repository not initialized' +_MISSING_MONGO = 'MongoDB repository not initialized' + +# Where a digest may sit inside the BFF-owned `request_data`. Checked in this order; absence is not +# a failure, only a disagreement is. +REQUEST_DIGEST_KEYS: Final[tuple[str, ...]] = ( + 'expected_digest', + 'digest', + 'sha256', + 'file_digest', +) + + +class ImportState(StrEnum): + """The workflow's internal vocabulary, mapped onto the shared status column by D17.""" + + RUNNING = 'RUNNING' + COMPLETED = 'COMPLETED' + ERROR = 'ERROR' + + +class ExperimentRunStatus(StrEnum): + """The values this workflow writes into `experiment_run.status`. + + Two are **reused** from the training flow because they describe a stage, not a flow — the row's + flow is told by `run_type`, not by the status. Two are **new in V12**, because + `TRAINING_SUCCESS` / `TRAINING_ERROR` would lie about where the row came from. Nobody should + "fix" the reuse into an `IMPORT_RUNNING` twin: it would carry no information `run_type` does not + already carry, and it would widen the pairing CHECK and the frontend's switch for nothing. + """ + + ORCHESTRATOR_WAITING_PROC = 'ORCHESTRATOR_WAITING_PROC' # reused + ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR' # reused + IMPORT_SUCCESS = 'IMPORT_SUCCESS' # new in V12 + IMPORT_ERROR = 'IMPORT_ERROR' # new in V12 + + +# A failure at these steps is the user's to act on — re-export, re-type the password, re-submit — +# which is the one thing the import screen can branch on without reading the error code. +USER_ACTIONABLE_STEPS: Final[frozenset[ImportStep]] = frozenset( + { + ImportStep.RECEIVED, + ImportStep.DOWNLOAD, + ImportStep.DECRYPTION, + ImportStep.ARCHIVE_INSPECTION, + ImportStep.EXTRACTION, + ImportStep.STRUCTURE_VALIDATION, + ImportStep.CONTENT_POLICY, + } +) + +# Only the status writes retry, and only because they are idempotent: `UPDATE ... SET +# WHERE id` leaves the row the same however many times it runs. A provisioning call has no such +# shape, which is why it keeps `maximum_attempts=1`. Defined here rather than in `sientia_do`, which +# other repositories share. +IMPORT_STATUS_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy( + maximum_attempts=10, + initial_interval=timedelta(seconds=1), + backoff_coefficient=2.0, + maximum_interval=timedelta(seconds=30), + non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'], +) + +# Progress writes are hints nobody reads — the terminal write rewrites them — so a 151-second pause +# before each of a dozen stage updates would stall a working import for twenty minutes because of a +# database the import itself does not need. +IMPORT_STATUS_HINT_RETRY_POLICY: Final[RetryPolicy] = RetryPolicy( + maximum_attempts=3, + initial_interval=timedelta(seconds=1), + backoff_coefficient=2.0, + maximum_interval=timedelta(seconds=2), + non_retryable_error_types=['ImportLogRowNotClaimableError', 'ImportStatusRejectedError'], +) + +_CLAIM_SQL = text( + """ + UPDATE public.experiment_run + SET status = :status, + updated_at = :updated_at, + orchestrator_response_data = CAST(:detail AS JSONB) + WHERE id = :import_run_id + AND run_type = 'IMPORT' + AND (status = 'PENDING' + OR (status = 'ORCHESTRATOR_WAITING_PROC' + AND orchestrator_response_data -> 'import' ->> 'workflow_id' = :workflow_id)) + """ +) + +_READ_ROW_SQL = text( + """ + SELECT id, + run_type, + status, + bucket_name, + file_name, + request_data, + orchestrator_response_data -> 'import' ->> 'workflow_id' AS claimed_by + FROM public.experiment_run + WHERE id = :import_run_id + """ +) + +# Class B and C key on the id alone: ownership was decided once, by the claim. Re-checking it on +# every later write buys nothing and costs the thing that matters — a predicate miss turning a +# pipeline that already imported the model into a reported failure. +_PROGRESS_SQL = text( + """ + UPDATE public.experiment_run + SET status = :status, + updated_at = :updated_at, + orchestrator_response_data = CAST(:detail AS JSONB) + WHERE id = :import_run_id + """ +) + +_NAMES_SQL = text( + """ + UPDATE public.experiment_run + SET experiment_name = :experiment_name, + run_name = :run_name, + updated_at = :updated_at, + orchestrator_response_data = CAST(:detail AS JSONB) + WHERE id = :import_run_id + """ +) + +_TERMINAL_SQL = text( + """ + UPDATE public.experiment_run + SET status = :status, + error_message = :error_message, + updated_at = :updated_at, + orchestrator_response_data = CAST(:detail AS JSONB) + WHERE id = :import_run_id + """ +) + + +def map_state_to_status( + state: ImportState, step: ImportStep | str | None = None +) -> ExperimentRunStatus: + """Map the workflow's internal state onto the shared status column (design D17). + + `RUNNING` is the orchestrator holding the row. `COMPLETED` is `IMPORT_SUCCESS`. An `ERROR` is + split in two by *who can act on it*: the user, for everything up to the content policy, or only + support, for everything after it — the same split the table already draws between + `ORCHESTRATOR_VALIDATION_ERROR` and `TRAINING_ERROR`. + """ + if state is ImportState.RUNNING: + return ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC + if state is ImportState.COMPLETED: + return ExperimentRunStatus.IMPORT_SUCCESS + if step is not None and ImportStep(step) in USER_ACTIONABLE_STEPS: + return ExperimentRunStatus.ORCHESTRATOR_VALIDATION_ERROR + return ExperimentRunStatus.IMPORT_ERROR + + +def derive_run_name(model_name: str, model_version: str) -> str: + """`import-{model_name}-v{model_version}`, truncated inside the model-name segment. + + Deterministic, because the recorded name and the created run's name must not diverge: the run is + created with exactly this value, passed in as an argument rather than recomputed. The `import-` + prefix guarantees the column's length-3 floor; the truncation is what keeps it under 50. + """ + suffix = f'-v{model_version}' + room = RUN_NAME_MAX_LENGTH - len(RUN_NAME_PREFIX) - len(suffix) + if room < 1: + # A version string long enough to leave no room for a name: keep the prefix and as much of + # the version as fits, so the value is still a legible, deterministic 50 characters. + return f'{RUN_NAME_PREFIX}{suffix.lstrip("-")}'[:RUN_NAME_MAX_LENGTH] + return f'{RUN_NAME_PREFIX}{model_name[:room]}{suffix}' + + +def build_import_detail( + *, + workflow_id: str, + step: ImportStep | str | None = None, + code: ImportErrorCode | str | None = None, + gate: int | None = None, + source: dict[str, Any] | None = None, + cleanup_failed: bool = False, +) -> str: + """Build `orchestrator_response_data` as JSON text. + + The whole `import` object is rewritten on every write, which is what keeps the workflow + identifier the claim wrote from ever being dropped. There is no column for any of this (design + D21): the code the frontend switches on and the identifier support quotes live at fixed keys in + the JSON column the workflow already owns. + """ + detail: dict[str, Any] = {'workflow_id': workflow_id, 'gate': gate} + if step is not None: + detail['step'] = str(step) + if code is not None: + detail['code'] = str(code) + if source: + detail['source'] = source + if cleanup_failed: + detail['cleanup_failed'] = True + return json.dumps({'import': detail}) + + +class ModelImport(SientiaMonitoring): + """Activities turning a `.sientia` object into a registered model, or into a durable ERROR.""" + + def __init__( + self, + status_config: dict[str, Any], + import_config: dict[str, Any], + minio_repository: MinioRepository | None = None, + mlflow_repository: MLFlowRepository | None = None, + mongo_repository: MongoDBRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + """ + Args: + status_config: `build_import_status_config()` — the BFF database holding the import log. + This is a **second** Postgres connection, independent of the one Laborious already + holds, which is what makes the transactional independence of design D7 literal. + import_config: `build_import_config()` — bucket, prefix, ceilings, work directory, + envelope key and the models collection name. + minio_repository: the shared MinIO repository, for the uploaded bundle. + mlflow_repository: the shared MLflow repository — the same instance the rest of the + worker uses, because setting the tracking URI is a process-wide side effect. + mongo_repository: MongoDB, where the model listing lives. + logger: logger instance. + notification_handler: the notification path the rest of the repository uses. + metrics_controller: metrics controller. + """ + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + self.import_config = import_config + self.status_config = status_config + self.import_minio_repository = minio_repository + self.import_mlflow_repository = mlflow_repository + self.mongo_repository = mongo_repository + self.status_engine = create_engine( + URL.create( + 'postgresql', + username=status_config['user'], + password=status_config['password'], + host=status_config['host'], + port=status_config['port'], + database=status_config['dbname'], + ), + pool_size=2, + max_overflow=3, + pool_pre_ping=True, + ) + + def close(self) -> None: + """Dispose the status engine. The import holds no other connection.""" + engine = getattr(self, 'status_engine', None) + if engine is not None: + engine.dispose() + + # ------------------------------------------------------------------ helpers + + @staticmethod + def _timestamp() -> Any: + """`updated_at` is `TIMESTAMP WITHOUT TIME ZONE`, so the value handed over is naive.""" + moment = now() + return moment.replace(tzinfo=None) if moment.tzinfo else moment + + @staticmethod + def _digest_from_request_data(request_data: Any) -> str | None: + """Find the uploader's digest inside the BFF-owned request payload, if it recorded one.""" + if not isinstance(request_data, dict): + return None + for key in REQUEST_DIGEST_KEYS: + value = request_data.get(key) + if isinstance(value, str) and value: + return value.lower() + return None + + def _execute_status_write(self, statement: Any, parameters: dict[str, Any]) -> int: + """Run one status statement in its own transaction and return the rows it affected. + + Raises: + ImportStatusRejectedError: the database accepted the connection and refused the + statement — a CHECK violation, an undefined column, a value too long. Non-retryable, + so a missing `V12` fails on the first attempt instead of after the full envelope. + Exception: anything else (connection refused, reset, timeout, a pod being replaced) is + left to propagate, because repeating it is exactly what fixes it. + """ + try: + with self.status_engine.begin() as connection: + return connection.execute(statement, parameters).rowcount + except (IntegrityError, ProgrammingError, DataError) as error: + raise ImportStatusRejectedError( + 'the import log database refused the status statement' + ) from error + + def _read_row(self, import_run_id: Any) -> dict[str, Any] | None: + """Read the record. A read, not a write: nothing of somebody else's is overwritten.""" + with self.status_engine.connect() as connection: + row = connection.execute(_READ_ROW_SQL, {'import_run_id': import_run_id}).mappings() + found = row.first() + return dict(found) if found is not None else None + + @staticmethod + def _describe_unclaimable(row: dict[str, Any] | None, workflow_id: str) -> str: + """Name the cause of a zero-row claim: missing row, training row, or someone else's.""" + if row is None: + return 'no such import log row' + if row.get('run_type') != 'IMPORT': + return f'row is a {row.get("run_type")} run, not an import' + claimed_by = row.get('claimed_by') + if claimed_by and claimed_by != workflow_id: + return f'row is already claimed by workflow {claimed_by}' + return f'row is in status {row.get("status")} and cannot be claimed' + + # -------------------------------------------------------------- class A: claim + + @activity.defn(name='claim_import_status') + async def claim_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Claim the import log row the frontend inserted. The workflow's first write. + + One conditional `UPDATE`, whose second disjunct makes a re-execution of the same workflow + idempotent instead of a conflict — a lost response on attempt 1 must not become a false + "not mine" on attempt 2. Zero rows affected covers the id typo, the duplicated start and the + wrong `run_type` in one condition, with no read-then-write race, and it is what stops a + duplicated workflow start from provisioning twice. + + Args (input_data): + metadata (dict): workflow metadata for logging and metrics. + import_run_id (int): the record's primary key, from the workflow input. + workflow_id (str): this workflow's identifier, recorded in the JSON detail. + source (dict, optional): the bundle's declared origin, when already known. + + Returns: + dict: `claimed`, plus the row values the input guards are checked against — + `bucket_name`, `file_name` and `request_digest`. + + Raises: + ImportLogRowNotClaimableError: the row is not this workflow's. Non-retryable, reported + through the channel, and **written to no row** — the loudness has to go somewhere + that is not somebody else's audit record. + """ + metadata = input_data.get('metadata', {}) + import_run_id = input_data['import_run_id'] + workflow_id = input_data['workflow_id'] + + affected = self._execute_status_write( + _CLAIM_SQL, + { + 'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value, + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=workflow_id, + step=ImportStep.RECEIVED, + source=input_data.get('source'), + ), + 'import_run_id': import_run_id, + 'workflow_id': workflow_id, + }, + ) + + row = self._read_row(import_run_id) + if affected == 0: + observed = self._describe_unclaimable(row, workflow_id) + error = ImportLogRowNotClaimableError( + f'import log row {import_run_id} could not be claimed: {observed}', + import_run_id=import_run_id, + observed=observed, + ) + await report_import_failure( + self, + None, + error, + import_run_id=import_run_id, + workflow_id=workflow_id, + metadata=metadata, + code=ImportErrorCode.LOG_ROW_NOT_CLAIMABLE, + marker='not_claimable', + ) + raise error + + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_STARTED_COUNT, + tags=self.get_core_labels(metadata, operation_type=ImportStep.RECEIVED.value), + ) + self.info(f'import log row {import_run_id} claimed by {workflow_id}', metadata) + row = row or {} + return { + 'claimed': True, + 'bucket_name': row.get('bucket_name'), + 'file_name': row.get('file_name'), + 'request_digest': self._digest_from_request_data(row.get('request_data')), + } + + # ----------------------------------------------------------- class B: progress + + @activity.defn(name='record_import_status') + async def record_import_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Record the step the workflow is about to attempt. A hint, keyed on the id alone. + + Its failure is **not** an import failure: the workflow reports it through the channel and + carries on, because losing a hint costs a record pointing one step behind and nothing else — + the terminal write rewrites status, code, reason and step from one classification. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail, so it is never dropped. + step (str): the step about to be attempted. + source (dict, optional): the bundle's declared origin. + """ + affected = self._execute_status_write( + _PROGRESS_SQL, + { + 'status': ExperimentRunStatus.ORCHESTRATOR_WAITING_PROC.value, + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=input_data['step'], + source=input_data.get('source'), + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + return {'recorded': affected > 0, 'rows': affected} + + @activity.defn(name='record_import_names') + async def record_import_names(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Write the experiment and run names, before anything exists in MLflow. + + Both columns exist so a person watching the import list sees *which* model is being + imported. The earliest moment both values are known is when the bundle is open, so that is + when they are written — before the experiment exists, before the run exists. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail. + experiment_name (str): the bundle's own experiment name, verbatim. + run_name (str): the derived run name, which the run is then created with. + source (dict, optional): the bundle's declared origin. + """ + affected = self._execute_status_write( + _NAMES_SQL, + { + 'experiment_name': input_data['experiment_name'], + 'run_name': input_data['run_name'], + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=input_data.get('step', ImportStep.EXPERIMENT_CREATION), + source=input_data.get('source'), + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + return {'recorded': affected > 0, 'rows': affected} + + # ----------------------------------------------------------- class C: terminal + + @activity.defn(name='record_import_terminal_status') + async def record_import_terminal_status(self, input_data: dict[str, Any]) -> dict[str, Any]: + """The single authoritative write, run from the workflow's `finally` after cleanup. + + Every failure is classified **once**, here, from the step that failed. On success the reason + column is written as `NULL` — never `''`, which its length-3 CHECK would refuse — and the + detail carries no code. A cleanup failure is noted in the detail without changing a + successful status: the user's model was imported, and reporting ERROR would be false. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): recorded in the detail. + succeeded (bool): whether every step succeeded. + step (str, optional): the failing step, or the last step reached on success. + gate (int, optional): the gate that rejected the bundle, when one did. + source (dict, optional): the bundle's declared origin. + cleanup_failed (bool, optional): whether the temporary-file cleanup failed. + + Returns: + dict: the verdict that was written — status, code and reason — so a caller that has to + report a failure to record it still holds the verdict it could not write. + """ + metadata = input_data.get('metadata', {}) + succeeded = bool(input_data.get('succeeded')) + step = input_data.get('step') + cleanup_failed = bool(input_data.get('cleanup_failed')) + + if succeeded: + status = map_state_to_status(ImportState.COMPLETED) + code: ImportErrorCode | None = None + reason: str | None = None + else: + failure = classify_import_failure(ImportStep(step) if step else None) + status = map_state_to_status(ImportState.ERROR, step) + code, reason = failure.code, failure.reason + + verdict = { + 'status': status.value, + 'code': code.value if code else None, + 'reason': reason, + 'step': str(step) if step else None, + } + + self._execute_status_write( + _TERMINAL_SQL, + { + 'status': status.value, + 'error_message': reason, + 'updated_at': self._timestamp(), + 'detail': build_import_detail( + workflow_id=input_data['workflow_id'], + step=step, + code=code, + gate=input_data.get('gate'), + source=input_data.get('source'), + cleanup_failed=cleanup_failed, + ), + 'import_run_id': input_data['import_run_id'], + }, + ) + + if succeeded: + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_COMPLETED_COUNT, + tags=self.get_core_labels(metadata, operation_type=status.value), + ) + self.info(f'import {input_data["import_run_id"]} recorded as {status.value}', metadata) + return verdict + + @activity.defn(name='report_import_status_write_failure') + async def report_import_status_write_failure( + self, input_data: dict[str, Any] + ) -> dict[str, Any]: + """Report a status write that could not be made — the one failure with no row to land on. + + Called by the workflow when a progress or terminal write exhausted its attempts or was + refused. It carries the verdict it could not write, so the outcome survives in text even + when it cannot survive in the row, and it never turns a completed import into a failed one. + + Args (input_data): + metadata (dict): workflow metadata. + import_run_id (int): the record's primary key. + workflow_id (str): the identifier that correlates this line with the record. + detail (str): the underlying failure, as text — it crossed an activity boundary, so the + exception object itself is gone by now. + verdict (str, optional): the verdict the terminal write could not record. + """ + await report_import_failure( + self, + None, + None, + import_run_id=input_data['import_run_id'], + workflow_id=input_data['workflow_id'], + metadata=input_data.get('metadata', {}), + code=ImportErrorCode.STATUS_NOT_RECORDED, + marker='status_write', + verdict=input_data.get('verdict'), + detail_text=input_data.get('detail'), + ) + return {'reported': True} + + # ------------------------------------------------------------- bundle helpers + + def _bundle_reader(self) -> BundleReader: + """A reader carrying the configured ceilings.""" + return BundleReader( + BundleLimits( + max_object_bytes=self.import_config['max_object_bytes'], + max_entries=self.import_config['max_entries'], + max_uncompressed_bytes=self.import_config['max_uncompressed_bytes'], + max_compression_ratio=self.import_config['max_compression_ratio'], + object_prefix=self.import_config['prefix'], + ) + ) + + def _work_dir(self, workflow_id: str) -> Path: + """This import's own scratch directory. One per workflow, never shared.""" + return Path(self.import_config['work_dir']) / workflow_id + + def _effective_bucket(self, requested: str | None) -> str: + return requested or self.import_config['bucket'] + + # --------------------------------------------------------- download and open + + @activity.defn(name='download_import_bundle') + async def download_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Gate 1, then the download, then the mandatory digest comparison. + + The object's size and key shape are checked from `stat_object`, before its body is fetched. + The SHA-256 of what was downloaded is then compared against `expected_digest` — the value the + uploader computed over the bytes it uploaded — with **no path that skips the comparison**, + because a missing value was already rejected at `received`. A mismatch fails here, before the + header is parsed: it means the object behind that key is not the file the user chose. + + Args (input_data): + metadata (dict): workflow metadata. + workflow_id (str): names this import's scratch directory. + bucket (str, optional): the effective bucket; the configured one when absent. + object_key (str): the uploaded object's key. + expected_digest (str): the uploader's SHA-256, lowercase hex. + + Returns: + dict: `path`, `size` and `digest` of the downloaded file. + + Raises: + ImportGateRejectedError: gate 1 rejected the object, or the digests disagree. + """ + if self.import_minio_repository is None: + raise ValueError(_MISSING_MINIO) + + metadata = input_data.get('metadata', {}) + bucket = self._effective_bucket(input_data.get('bucket')) + object_key = input_data['object_key'] + expected_digest = input_data['expected_digest'] + reader = self._bundle_reader() + + client = self.import_minio_repository.minio_client + try: + stat = client.stat_object(bucket, object_key) + except Exception as error: + raise ImportGateRejectedError( + 'the uploaded object could not be read', step=ImportStep.DOWNLOAD, gate=1 + ) from error + + try: + reader.check_object_limits(size=stat.size or 0, key=object_key) + except BundleGateError as error: + raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error + + work_dir = self._work_dir(input_data['workflow_id']) + work_dir.mkdir(parents=True, exist_ok=True) + target = work_dir / 'bundle.sientia' + try: + client.fget_object(bucket, object_key, str(target)) + except Exception as error: + raise ImportGateRejectedError( + 'the uploaded object could not be downloaded', step=ImportStep.DOWNLOAD, gate=1 + ) from error + + digest = fmt.file_digest(target) + if digest != str(expected_digest).lower(): + target.unlink(missing_ok=True) + raise ImportGateRejectedError( + 'the downloaded object does not match the digest the uploader computed', + step=ImportStep.DOWNLOAD, + gate=1, + ) + + self.info(f'import bundle downloaded from {bucket} ({stat.size} bytes)', metadata) + return {'path': str(target), 'size': stat.size, 'digest': digest} + + @activity.defn(name='open_import_bundle') + async def open_import_bundle(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Open the envelope, then the bundle: gates 2 to 7, in the reader's own order. + + The password envelope is decrypted **here**, inside the activity that uses the password, so + no plaintext password is ever an activity input, an activity result or a workflow variable — + and therefore never lands in the Temporal event history. The recovered password is a local + value and nothing else. + + Args (input_data): + metadata (dict): workflow metadata. + workflow_id (str): names this import's scratch directory. + path (str): the downloaded bundle. + password_envelope (str): base64 `nonce || ciphertext || tag`, AES-256-GCM. + + Returns: + dict: `extracted_dir`, `metadata`, `parameters`, `digest` — nothing secret. + + Raises: + ImportGateRejectedError: with the step and the gate that rejected the bundle. + """ + metadata = input_data.get('metadata', {}) + work_dir = self._work_dir(input_data['workflow_id']) + extract_root = work_dir / 'extract' + + try: + key = load_envelope_key(self.import_config.get('password_key')) + password = decrypt_password_envelope(input_data['password_envelope'], key) + except EnvelopeError as error: + # A distinct failure from a wrong *bundle* password, and never confused with it: the + # envelope is the platform's own transport, and its message says so. Same step, because + # the step is where the import stopped, not who is at fault. + raise ImportGateRejectedError( + 'the import password envelope could not be opened', + step=ImportStep.DECRYPTION, + gate=3, + cause='envelope', + ) from error + + reader = self._bundle_reader() + try: + opened = reader.open(input_data['path'], password, extract_root, work_dir) + except BundleGateError as error: + raise ImportGateRejectedError(str(error), step=error.step, gate=error.gate) from error + finally: + del password + + await self.emit_metric( + metric_object=metrics.MODEL_IMPORT_KDF_DURATION, + tags=self.get_core_labels(metadata, operation_type=ImportStep.DECRYPTION.value), + method='observe', + value=opened.kdf_seconds, + ) + self.info(f'import bundle opened into {opened.extracted_dir}', metadata) + return { + 'extracted_dir': str(opened.extracted_dir), + 'metadata': opened.metadata, + 'parameters': opened.parameters, + 'digest': opened.digest, + } + + # ------------------------------------------------------------ MLflow provisioning + + @activity.defn(name='create_import_experiment') + async def create_import_experiment(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Get or create the experiment named by the bundle, keeping the origin's name. + + Args (input_data): + metadata (dict): workflow metadata. + experiment_name (str): the bundle's `metadata.experiment_name`. + + Returns: + dict: `experiment_id` and `experiment_name`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + experiment_name = input_data['experiment_name'] + experiment = self.import_mlflow_repository.get_experiment( + experiment_name, create_if_not_exists=True + ) + self.info( + f'import experiment ready: {experiment_name} ({experiment.experiment_id})', + input_data.get('metadata', {}), + ) + return {'experiment_id': experiment.experiment_id, 'experiment_name': experiment_name} + + @activity.defn(name='upload_import_artifacts') + async def upload_import_artifacts(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Create a new run with the recorded name and log the bundle's tree and parameters. + + Always a new run — the importer never appends to an existing one. The run is created with + exactly the `run_name` already written on the record, passed in rather than recomputed, so + the two cannot diverge. The artifact tree is logged verbatim, keeping `prediction_model` and + `data_model` at the run's artifact root, which is what the legacy loader resolves. The + bundle's `parameters` are logged unchanged, and the origin's own identity is logged + alongside them under `source.*` — records of where the model came from, not links. + + Args (input_data): + metadata (dict): workflow metadata. + experiment_id (str): the experiment to create the run in. + run_name (str): the name already on the record. + extracted_dir (str): the opened bundle's directory. + parameters (dict): the bundle's `parameters` block, verbatim. + source (dict): the bundle's `metadata` block, logged as `source.*`. + + Returns: + dict: `run_id` and `run_name`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + client = self.import_mlflow_repository.client + run = client.create_run( + experiment_id=input_data['experiment_id'], run_name=input_data['run_name'] + ) + run_id = run.info.run_id + + artifacts_dir = Path(input_data['extracted_dir']) / 'artifacts' + client.log_artifacts(run_id, str(artifacts_dir)) + + for name, value in (input_data.get('parameters') or {}).items(): + client.log_param(run_id, name, value) + for name, value in (input_data.get('source') or {}).items(): + client.log_param(run_id, f'source.{name}', value) + + client.set_terminated(run_id, 'FINISHED') + self.info(f'import artifacts uploaded to run {run_id}', input_data.get('metadata', {})) + return {'run_id': run_id, 'run_name': input_data['run_name']} + + @activity.defn(name='register_import_model_version') + async def register_import_model_version(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Register a version whose source the legacy loader can parse. No stage change. + + The source is `runs://prediction_model`, because `get_model_run_id` extracts the run + id as `source.split('/')[2]` — a `models:/m-*` source, what the new stack emits, breaks it. + The stage is left untouched: a `COMPLETED` import is a registered model, not a servable one, + and promoting it to `Production` is a manual operator step outside this workflow. + + Args (input_data): + metadata (dict): workflow metadata. + model_name (str): the bundle's `metadata.model_name`, verbatim. + run_id (str): the run this import created. + + Returns: + dict: `model_name`, `version` and `source`. + """ + if self.import_mlflow_repository is None: + raise ValueError(_MISSING_MLFLOW) + + client = self.import_mlflow_repository.client + model_name = input_data['model_name'] + run_id = input_data['run_id'] + + try: + client.get_registered_model(model_name) + except Exception: + # First import of this name: the registered model is created on demand. A second import + # lands as a new version of the same registered model, never a differently named one. + client.create_registered_model(model_name) + + source = f'runs:/{run_id}/prediction_model' + version = client.create_model_version(name=model_name, source=source, run_id=run_id) + self.info( + f'import registered {model_name} version {version.version}', + input_data.get('metadata', {}), + ) + return {'model_name': model_name, 'version': str(version.version), 'source': source} + + # --------------------------------------------------- the MongoDB model document + + @activity.defn(name='write_import_model_document') + async def write_import_model_document(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Write the document that *is* the model listing, inert, last. + + This is what the runtime reads to know how to load a model — `model_config.transform_flavor`, + `model_config.predict_flavor` and `model_config.target` — so without it an imported model is + invisible to every workflow. `id` continues the collection's own numbering: the largest + existing id **as an integer** plus one, never a count and never a lexicographic sort, which + would answer `"9" > "30"`. A name already in the collection is left exactly as it is: + `active`, `target` and the flavors of a model already on the platform are operator-owned. + + Args (input_data): + metadata (dict): workflow metadata. + model_name (str): the bundle's `metadata.model_name`, verbatim. + target (str): the bundle's `parameters.target_variable`. + + Returns: + dict: `written`, `id`, `skipped` and how many unusable ids were ignored. + """ + if self.mongo_repository is None: + raise ValueError(_MISSING_MONGO) + + metadata = input_data.get('metadata', {}) + collection = self.import_config['models_collection'] + model_name = input_data['model_name'] + + existing = await self.mongo_repository.find_one(collection, {'name': model_name}, metadata) + if existing is not None: + self.info( + f'model document for {model_name} already exists; leaving it untouched', metadata + ) + return {'written': False, 'skipped': True, 'id': existing.get('id')} + + next_id, ignored = await self._next_model_document_id(collection, metadata) + document = { + 'id': next_id, + 'name': model_name, + 'active': False, + 'model_config': { + 'transform_flavor': 'sklearn', + 'predict_flavor': 'joblib', + 'retention_minutes': 60, + 'target': input_data['target'], + }, + } + await self.mongo_repository.insert(collection, document, metadata) + self.info(f'model document written for {model_name} with id {next_id}', metadata) + return {'written': True, 'skipped': False, 'id': next_id, 'ignored_ids': ignored} + + async def _next_model_document_id( + self, collection: str, metadata: dict[str, Any] + ) -> tuple[str, int]: + """`max(int(id)) + 1` as a decimal string, tolerating ids that are not numbers.""" + if self.mongo_repository is None: + raise ValueError(_MISSING_MONGO) + raw_ids = await self.mongo_repository.distinct(collection, 'id', {}, metadata) + numeric: list[int] = [] + ignored = 0 + for value in raw_ids: + try: + numeric.append(int(str(value))) + except (TypeError, ValueError): + # A document whose id is not a decimal number cannot take part in the numbering. + # Ignored rather than fatal, and counted so the log says how many there were. + ignored += 1 + if ignored: + self.info(f'ignored {ignored} non-numeric model document ids', metadata) + return str(max(numeric, default=0) + 1), ignored + + # ----------------------------------------------------------- cleanup and retention + + @activity.defn(name='cleanup_import_files') + async def cleanup_import_files(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Remove this import's scratch directory. Never the uploaded object. + + Runs on every path, success or failure, and **never raises**: a cleanup failure is reported + through the channel and noted in the record's detail, but the user's model was imported and + reporting an error would be false. + + Args (input_data): + metadata (dict): workflow metadata. + workflow_id (str): names the directory to remove. + import_run_id (int): for correlation when the removal fails. + + Returns: + dict: `cleaned` and, when it failed, the reason as text. + """ + metadata = input_data.get('metadata', {}) + work_dir = self._work_dir(input_data['workflow_id']) + try: + shutil.rmtree(work_dir, ignore_errors=False) + except FileNotFoundError: + # Nothing was downloaded, or a gate removed its own partial tree. Not a failure. + return {'cleaned': True, 'existed': False} + except OSError as error: + await report_import_failure( + self, + ImportStep.CLEANUP, + error, + import_run_id=input_data.get('import_run_id'), + workflow_id=input_data['workflow_id'], + metadata=metadata, + code=ImportErrorCode.CLEANUP_INCOMPLETE, + marker=ImportStep.CLEANUP.value, + ) + return {'cleaned': False, 'existed': True, 'reason': type(error).__name__} + return {'cleaned': True, 'existed': True} + + @activity.defn(name='ensure_import_bundle_retention') + async def ensure_import_bundle_retention(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Ensure the 7-day expiry rule on the import prefix, idempotently. + + Applied by the importer rather than assumed from a Helm value, so the acceptance criterion + ("lifecycle active") is verifiable from the code that depends on it. The uploaded object is + never deleted by the import itself — it is left to this rule, as the record of an attempt. + + Args (input_data): + metadata (dict): workflow metadata. + bucket (str, optional): the effective bucket; the configured one when absent. + + Returns: + dict: whether the rule was already there and what it expires. + """ + if self.import_minio_repository is None: + raise ValueError(_MISSING_MINIO) + + metadata = input_data.get('metadata', {}) + bucket = self._effective_bucket(input_data.get('bucket')) + prefix = self.import_config['prefix'] + days = self.import_config['retention_days'] + rule_id = 'sientia-import-bundle-expiry' + client = self.import_minio_repository.minio_client + + try: + existing = client.get_bucket_lifecycle(bucket) + except Exception: + # A bucket with no lifecycle configuration at all answers with an error on some + # deployments and with `None` on others; both mean "there is no rule yet". + existing = None + + rules = [rule for rule in (existing.rules if existing else []) if rule.rule_id != rule_id] + rules.append( + Rule( + ENABLED, + rule_filter=Filter(prefix=prefix), + rule_id=rule_id, + expiration=Expiration(days=days), + ) + ) + client.set_bucket_lifecycle(bucket, LifecycleConfig(rules)) + self.info(f'import retention rule ensured on {bucket} ({days} days)', metadata) + return {'bucket': bucket, 'prefix': prefix, 'days': days, 'rules': len(rules)} diff --git a/laborious/activities/model_import_errors.py b/laborious/activities/model_import_errors.py new file mode 100644 index 0000000..657f47f --- /dev/null +++ b/laborious/activities/model_import_errors.py @@ -0,0 +1,348 @@ +"""The import's failure vocabulary: a closed catalog, typed errors and one reporting channel. + +Three things live here and nowhere else: + +- **the catalog** — one stable `error_code` and one plain-language `error_reason` per pipeline step. + The reason is written for a person who does not know this system exists: what happened, what state + the model is in, what to do next. It never carries a path, a table name, an object key, a Python + exception class or a traceback, and it is never built from `str(exception)`. +- **the typed errors** — the two failures that must not be retried, because repeating them cannot + change the answer: a record this workflow does not own, and a statement the database refuses. +- **the channel** — `report_import_failure`, the single funnel every failure passes through. It is + the only place in the import path that logs an error or sends a notification, which a test + asserts, so no failure is reported twice or half-reported. +""" + +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + import traceback + from collections.abc import Mapping + from dataclasses import dataclass + from enum import StrEnum + from types import MappingProxyType + from typing import Any, Final, Protocol + + from temporalio.exceptions import ApplicationError + + from laborious import metrics + from laborious.utils.bundle.steps import ImportStep + + +class ImportErrorCode(StrEnum): + """The machine-readable half of a failure: the frontend's i18n key and support's triage key. + + Written under `import.code` in the record's JSON detail — not in a column of its own, so the + import adds nothing to a table it shares with the training flow. + """ + + REQUEST_INVALID = 'IMPORT_REQUEST_INVALID' + FILE_UNREADABLE = 'IMPORT_FILE_UNREADABLE' + BUNDLE_CANNOT_BE_OPENED = 'IMPORT_BUNDLE_CANNOT_BE_OPENED' + BUNDLE_REJECTED = 'IMPORT_BUNDLE_REJECTED' + BUNDLE_UNPACK_FAILED = 'IMPORT_BUNDLE_UNPACK_FAILED' + BUNDLE_INCOMPLETE = 'IMPORT_BUNDLE_INCOMPLETE' + BUNDLE_UNEXPECTED_CONTENT = 'IMPORT_BUNDLE_UNEXPECTED_CONTENT' + STORAGE_PREPARATION_FAILED = 'IMPORT_STORAGE_PREPARATION_FAILED' + MODEL_FILES_NOT_STORED = 'IMPORT_MODEL_FILES_NOT_STORED' + MODEL_NOT_REGISTERED = 'IMPORT_MODEL_NOT_REGISTERED' + MODEL_SETTINGS_NOT_SAVED = 'IMPORT_MODEL_SETTINGS_NOT_SAVED' + UNEXPECTED_ERROR = 'IMPORT_UNEXPECTED_ERROR' + # Codes that never reach a record's `error_message`, and are therefore absent from the catalog: + # one names the failure to own a row, one the failure to write one, one the temporary files left + # on the worker after an import that otherwise succeeded. All three live on the channel — log, + # metric, notification, Temporal error — and the third is additionally flagged in the record's + # detail as `cleanup_failed`, without a code and without changing the status. + LOG_ROW_NOT_CLAIMABLE = 'IMPORT_LOG_ROW_NOT_CLAIMABLE' + STATUS_NOT_RECORDED = 'IMPORT_STATUS_NOT_RECORDED' + CLEANUP_INCOMPLETE = 'IMPORT_CLEANUP_INCOMPLETE' + + +@dataclass(frozen=True) +class ImportFailure: + """One classification of one failure: the step, its code and the sentence it stores.""" + + step: ImportStep | None + code: ImportErrorCode + reason: str + + +_CONTACT_SUPPORT = "Contact support with this import's identifier." +_RE_EXPORT = 'Export the model again from the origin platform.' + +IMPORT_FAILURE_CATALOG: Final[Mapping[ImportStep, tuple[ImportErrorCode, str]]] = MappingProxyType( + { + ImportStep.RECEIVED: ( + ImportErrorCode.REQUEST_INVALID, + 'The import could not be started because the request was incomplete or ' + 'contradictory. Nothing was created. Submit the import again.', + ), + ImportStep.DOWNLOAD: ( + ImportErrorCode.FILE_UNREADABLE, + 'The uploaded file could not be read, or it is not the file that was sent. ' + 'Upload it again.', + ), + ImportStep.DECRYPTION: ( + ImportErrorCode.BUNDLE_CANNOT_BE_OPENED, + 'This file could not be opened: the password is wrong, or the file is damaged or ' + 'was altered. Confirm the password with whoever exported the model, then try ' + 'again.', + ), + ImportStep.ARCHIVE_INSPECTION: ( + ImportErrorCode.BUNDLE_REJECTED, + "The file's contents did not pass the platform's safety checks, so it was not " + f'opened. {_RE_EXPORT}', + ), + ImportStep.EXTRACTION: ( + ImportErrorCode.BUNDLE_UNPACK_FAILED, + f'The file could not be unpacked safely and nothing from it was kept. {_RE_EXPORT}', + ), + ImportStep.STRUCTURE_VALIDATION: ( + ImportErrorCode.BUNDLE_INCOMPLETE, + 'This file is not a complete model export — part of what the platform needs is ' + f'missing from it. {_RE_EXPORT}', + ), + ImportStep.CONTENT_POLICY: ( + ImportErrorCode.BUNDLE_UNEXPECTED_CONTENT, + 'The file contains something a model export should not contain, so it was ' + f'rejected. {_RE_EXPORT}', + ), + ImportStep.EXPERIMENT_CREATION: ( + ImportErrorCode.STORAGE_PREPARATION_FAILED, + 'The platform could not prepare a place to keep this model. The import stopped ' + f'and no model was created. {_CONTACT_SUPPORT}', + ), + ImportStep.ARTIFACT_UPLOAD: ( + ImportErrorCode.MODEL_FILES_NOT_STORED, + "The model's files could not be stored on this platform. The import stopped and " + f'the model is not available. {_CONTACT_SUPPORT}', + ), + ImportStep.REGISTRATION: ( + ImportErrorCode.MODEL_NOT_REGISTERED, + "The model's files were stored, but the model itself could not be registered, so " + f'it cannot be used. {_CONTACT_SUPPORT}', + ), + ImportStep.MODEL_DOCUMENT: ( + ImportErrorCode.MODEL_SETTINGS_NOT_SAVED, + 'The model was registered, but the settings that tell the platform how to run it ' + f'could not be saved. {_CONTACT_SUPPORT}', + ), + } +) + +UNCLASSIFIED_FAILURE: Final[tuple[ImportErrorCode, str]] = ( + ImportErrorCode.UNEXPECTED_ERROR, + f'The import stopped for an unexpected reason and no model was created. {_CONTACT_SUPPORT}', +) + +# `cleanup` deletes temporary files on the worker. It is never the reported failure: when every +# other step succeeded and only cleanup failed, the import is a success and the leak is a log line. +# So it deliberately has no user-facing sentence, and the catalog-completeness test knows it. +STEPS_WITHOUT_A_USER_FACING_REASON: Final[frozenset[ImportStep]] = frozenset({ImportStep.CLEANUP}) + + +def classify_import_failure( + step: ImportStep | None, exception: BaseException | None = None +) -> ImportFailure: + """Turn a failure into the one code and the one sentence that will be recorded. + + The exception's text is **ignored entirely** — it is what carries paths, table names and type + names, and it belongs in the log. What decides the outcome is the step. + + Args: + step: the step that was being attempted, or `None` when none had been entered. + exception: the failure, when there is one to hand. It is accepted so callers read as + "classify this failure" and never absent because the classification needs it: the + catalog is keyed by step, and an exception cannot cross an activity boundary with its + type intact anyway — the terminal write classifies from the step alone. + + Returns: + ImportFailure: step, code and the reason to store. + """ + del exception # the catalog is keyed by step; the exception's text never reaches the record + if step is None or step not in IMPORT_FAILURE_CATALOG: + code, reason = UNCLASSIFIED_FAILURE + return ImportFailure(step=step, code=code, reason=reason) + code, reason = IMPORT_FAILURE_CATALOG[step] + return ImportFailure(step=step, code=code, reason=reason) + + +class ImportLogRowNotClaimableError(ApplicationError): + """The record named by `import_run_id` is not this workflow's to write. + + Raised when the claim affects zero rows: no such row, a `TRAIN` row, or a row already running + under another workflow. Non-retryable on purpose — a second attempt would read the same rows and + reach the same conclusion — and loud everywhere except the row itself, which belongs to someone + else. + """ + + def __init__(self, message: str, *, import_run_id: Any = None, observed: str = '') -> None: + super().__init__( + message, + type='ImportLogRowNotClaimableError', + non_retryable=True, + ) + self.import_run_id = import_run_id + self.observed = observed + + +class ImportStatusRejectedError(ApplicationError): + """The database accepted the connection and refused the statement. + + A `CHECK` violation, an undefined column, a value too long for its column. `V12` not being + applied yet lands here, at the terminal write, and it must fail on the first attempt instead of + spending the whole retry envelope to say so. + """ + + def __init__(self, message: str) -> None: + super().__init__(message, type='ImportStatusRejectedError', non_retryable=True) + + +class ImportInputError(ApplicationError): + """The request is incomplete or contradictory. Fails at `received`, on the claimed row.""" + + def __init__(self, message: str) -> None: + super().__init__( + message, + {'step': ImportStep.RECEIVED.value, 'gate': None}, + type='ImportInputError', + non_retryable=True, + ) + + +class ImportGateRejectedError(ApplicationError): + """A gate rejected the bundle, and the error says which one. + + Gates 2 to 7 all run inside one activity, so the step the workflow was *attempting* is not + precise enough to record. An exception loses its Python type crossing an activity boundary but + keeps its `details`, so the step and the gate travel there and the terminal write can name them. + """ + + def __init__( + self, + message: str, + *, + step: ImportStep, + gate: int | None = None, + cause: str | None = None, + ) -> None: + super().__init__( + message, + {'step': step.value, 'gate': gate, 'cause': cause}, + type='ImportGateRejectedError', + non_retryable=True, + ) + self.step = step + self.gate = gate + + +def failure_location(exception: BaseException | None) -> tuple[str | None, int | None]: + """Read the step and gate an activity's failure carries, if it carried any. + + Walks the `__cause__` chain because Temporal wraps an activity failure in an `ActivityError`. + Returns `(None, None)` when the failure did not say where it happened, and the caller then falls + back to the step it was attempting. + """ + seen = 0 + candidate: BaseException | None = exception + while candidate is not None and seen < 10: + details = getattr(candidate, 'details', None) + if details and isinstance(details[0], dict) and 'step' in details[0]: + location = details[0] + gate = location.get('gate') + return location.get('step'), int(gate) if isinstance(gate, int) else None + candidate = candidate.__cause__ + seen += 1 + return None, None + + +class FailureSinks(Protocol): + """What the channel needs from its caller: the repository's existing observability surface.""" + + def error(self, message: str, metadata: dict | None = None) -> None: ... + + def get_core_labels(self, metadata: dict | None = None, operation_type: str = '-') -> dict: ... + + async def emit_metric( + self, + metric_object: Any, + tags: dict[str, Any], + method: str = 'inc', + value: float | None = None, + ) -> None: ... + + async def send_notification_async( + self, + metadata: dict, + notification_id: str, + message: str, + block: str, + level: str = 'INFO', + attachment_content: str | None = None, + ) -> None: ... + + +async def report_import_failure( + sinks: FailureSinks, + step: ImportStep | None, + exception: BaseException | None, + *, + import_run_id: Any, + workflow_id: str, + metadata: dict[str, Any] | None = None, + code: ImportErrorCode | None = None, + marker: str | None = None, + verdict: str | None = None, + detail_text: str | None = None, +) -> ImportFailure: + """Report one failure through every sink, exactly once. + + Sinks, in order: a structured log carrying the ids, the step, the code **and** the exception + with its traceback (the only sink that gets the exception text); the import error metric, with + the code in `operation_type`; a notification through the existing handler. The exception itself + reaches Temporal by being raised, and the record is the fifth, optional sink — the terminal + write, when the row can be written at all. + + Args: + sinks: the activity providing logger, metrics and notifications. + step: the pipeline step, or `None` for a failure that is not one (see `marker`). + exception: the failure. + import_run_id: the record's id, for correlation. + workflow_id: the identifier a user quotes to support. + metadata: workflow metadata for labels and the notification. + code: an explicit code, for the two failures that are not pipeline steps. + marker: `not_claimable` or `status_write` — what a step name would have been. + verdict: for a failed terminal write, the verdict it could not record, so the outcome + survives in text even when it cannot survive in the row. + detail_text: the failure as text, for the callers whose exception object is already gone — + anything that crossed an activity boundary. Used in place of a traceback. + + Returns: + ImportFailure: the classification that was reported. + """ + failure = classify_import_failure(step, exception) + if code is not None: + failure = ImportFailure(step=step, code=code, reason=failure.reason) + + where = marker or (step.value if step else 'unknown') + if exception is not None: + detail = ''.join(traceback.format_exception(exception)).strip() + else: + detail = detail_text or 'no exception object available' + verdict_line = f' verdict_not_recorded={verdict}' if verdict else '' + sinks.error( + f'model import failed: step={where} code={failure.code.value} ' + f'import_run_id={import_run_id} workflow_id={workflow_id}{verdict_line}\n{detail}', + metadata or {}, + ) + + labels = sinks.get_core_labels(metadata or {}, operation_type=failure.code.value) + await sinks.emit_metric(metric_object=metrics.MODEL_IMPORT_ERROR_COUNT, tags=labels) + await sinks.send_notification_async( + metadata=metadata or {}, + notification_id=failure.code.value, + message=failure.reason, + block=where, + level='ERROR', + ) + return failure diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py new file mode 100644 index 0000000..79b140c --- /dev/null +++ b/laborious/activities/model_metrics.py @@ -0,0 +1,364 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import time + import traceback + import warnings + from typing import Any + + import numpy as np + from pandas import DataFrame, Index, to_datetime + from sientia.ModelAnalysis import ModelAnalysis + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ + + from laborious import metrics + from laborious.utils.dataframe_debug import build_dataframe_debug_message + +warnings.filterwarnings('ignore', category=RuntimeWarning, message='Degrees of freedom <= 0') +warnings.filterwarnings( + 'ignore', category=RuntimeWarning, message='invalid value encountered in scalar divide' +) + + +class ModelMetrics(SientiaMonitoring): + """ + Metrics activities for the Laborious system. + + This class provides activities for writing metrics to the Prometheus monitoring system. + """ + + _MAX_DEBUG_DATAFRAME_ROWS = 100 + + def __init__( + self, + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + ): + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + def close(self) -> None: + """ + Close the model metrics activity and clean up resources. + """ + SientiaMonitoring.shutdown(self) + + def __del__(self): + self.close() + + def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None: + """ + Log dataframe content only when row count is below the configured threshold + + Args: + - message (str): Base log message to identify the dataframe in logs + - data (Any): Dataframe-like payload to be logged + - metadata (dict[str, Any]): Workflow metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + + async def get_drift_metrics( + self, + reference_data: DataFrame, + target_data: DataFrame, + target_name: str, + reference_columns: Index, + drift_metrics: list[str], + chunk_period: str, + metadata: dict[str, Any], + ) -> DataFrame: + """ + Calculate univariate drift metrics for a model. + Args: + model_analysis (ModelAnalysis): Model analysis object + reference_data (DataFrame): Reference data + target_data (DataFrame): Target data + reference_columns (list[str]): Reference columns + drift_metrics (list[str]): Drift metrics + metadata (dict[str, Any]): Workflow execution metadata + """ + + config = { + 'target': target_name, + 'prediction': 'prediction', + 'timestamp': 'timestamp', + 'features': reference_columns, + } + + model_analysis = ModelAnalysis(config=config) + + self._debug_dataframe( + f'Reference data: Size {reference_data.shape}', reference_data, metadata + ) + + self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata) + + core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift') + start_time = time.time() + try: + univariate_drift = model_analysis.detect_univariate_drift( + reference_df=reference_data, + analysis_df=target_data, + features=reference_columns, + timestamp_col=config['timestamp'], + methods=drift_metrics, + chunk_period=chunk_period, + ) + except Exception as e: + self.error(f'Error detecting univariate drift: {e}', metadata) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + core_labels = self.get_core_labels(metadata, operation_type='detect_multivariate_drift') + start_time = time.time() + try: + multivariate_drift = model_analysis.detect_multivariate_drift( + reference_df=reference_data, + analysis_df=target_data, + features=reference_columns, + timestamp_col=config['timestamp'], + chunk_period=chunk_period, + ) + except Exception as e: + self.error(f'Error detecting multivariate drift: {e}', metadata) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + start_time = time.time() + core_labels = self.get_core_labels(metadata, operation_type='get_drift_metrics_dataframe') + try: + drift_df = model_analysis.get_drift_metrics_dataframe( + univariate_drift=univariate_drift, + multivariate_drift=multivariate_drift, + ) + except Exception as e: + self.error(f'Error getting drift metrics: {e}', metadata) + await self.emit_metric( + metric_object=metrics.MODEL_ANALYZE_ERROR_COUNT, tags=core_labels + ) + raise e + await self.observe_lag(start_time, metrics.MODEL_ANALYZE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_ANALYZE_COUNT, tags=core_labels) + + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) + + return drift_df + + @activity.defn(name='calculate_drift') + async def calculate_drift(self, input_data: dict[str, Any]) -> list[dict]: + """ + Calculate drift metrics for a model. + + Args: + input_data (dict[str, Any]): Input data containing: + - metadata (dict): Workflow execution metadata + - model_name (str): Name of the MLFlow model to calculate drift for + - reference_data (pd.DataFrame): Reference data for the model + - target_data (pd.DataFrame): Target data for calculating drift + - target_name (str): Name of the target column + - drift_metrics (list[str]): List of drift metrics to calculate + """ + metadata = input_data['metadata'] + model_name = input_data['model_name'] + model_id = input_data['model_id'] + reference_raw_data = input_data['reference_data'] + target_data = DataFrame(input_data['target_data']) + target_name = input_data['target_name'] + drift_metrics = input_data['drift_metrics'] + chunk_period = input_data['chunk_period'] + + if chunk_period not in ['min', 's']: + self.error(f'Invalid chunk period: {chunk_period}', metadata) + raise ValueError(f'Invalid chunk period: {chunk_period}, must be "min" or "s"') + + self.info(f'Calculating drift for model {model_name}', metadata) + + target_data = target_data.pivot(index='timestamp', columns='variable', values='value') + target_data['timestamp'] = target_data.index + target_data['timestamp'] = to_datetime(target_data['timestamp']) + target_data['timestamp'] = target_data['timestamp'].dt.strftime(DATETIME_FORMAT) + target_data = target_data.reset_index(drop=True) + target_data.dropna(inplace=True) + + if reference_raw_data is not None: + self.info('Using reference data', metadata) + reference_data = DataFrame(reference_raw_data) + accurate = True + else: + # Get 30% first rows of target_data + self.warning('Using 30% first rows of target data as reference data', metadata) + target_data.sort_values(by='timestamp', ascending=True, inplace=True) + reference_data = target_data.head(int(len(target_data) * 0.3)) + accurate = False + + await self.send_notification_async( + metadata=metadata, + notification_id='MODEL_METRICS_REFERENCE_DATA_WARNING', + message='Using 30% first rows of target data as reference data', + block='model_metrics', + level=NotificationLevel.WARNING, + attachment_content=reference_data.to_csv(), + ) + + reference_columns = reference_data.drop( + columns=[target_name, 'timestamp', 'target', 'prediction'], errors='ignore' + ).columns + + try: + drift_df = await self.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name=target_name, + reference_columns=reference_columns, + drift_metrics=drift_metrics, + chunk_period=chunk_period, + metadata=metadata, + ) + except Exception as e: + self.error(f'Error getting drift metrics: {e}', metadata) + await self.send_notification_async( + metadata=metadata, + notification_id='MODEL_METRICS_GET_DRIFT_METRICS_ERROR', + message=f'Error getting drift metrics: {e}', + block='model_metrics', + level=NotificationLevel.ERROR, + attachment_content=traceback.format_exc(), + ) + return [] + + if drift_df.empty: + self.warning('No drift metrics found', metadata) + return [] + + # Drop unnecessary columns + drift_df.drop(columns=['p_value'], inplace=True) + + # Extract timestamps only until minutes + if chunk_period == 'min': + target_timestamps = target_data['timestamp'].apply(lambda x: x[:16]) + else: + target_timestamps = target_data['timestamp'] + + # Drop rows where timestamp is not in target data, to avoid save drift from reference + drift_df = drift_df[drift_df['timestamp'].isin(target_timestamps)] + + if drift_df.empty: + self.warning( + 'No drift metrics found after dropping rows where timestamp is not in target data', + metadata, + ) + return [] + + # Rename columns to match database columns + drift_df.rename( + columns={ + 'metric': 'method', + 'statistic': 'value', + }, + inplace=True, + ) + + # Drop duplicates + drift_df.drop_duplicates( + subset=['timestamp', 'method', 'feature'], keep='first', inplace=True + ) + + drift_df['model_id'] = model_id + drift_df['accurate'] = accurate + + drift_df['timestamp'] = to_datetime(drift_df['timestamp']) + drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC') + drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ) + + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) + + return drift_df.to_dict(orient='records') + + @activity.defn(name='calculate_simple_metrics') + async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> list[dict]: + """ + Calculate simple metrics for a model. Metrics available are: + - rmse + - mse + - mae + - r2 + - accuracy + - precision + - recall + - f1 + Args: + input_data (dict[str, Any]): Input data containing: + - metadata (dict): Workflow execution metadata + - model_id (str): ID of the MLFlow model + - target_data (pd.DataFrame): Target data for calculating metrics, containing target and prediction columns + - metrics (list[str]): List of metrics to calculate + Returns: + dict[Hashable, Any]: Dictionary containing the calculated metrics + """ + + metadata = input_data['metadata'] + model_id = input_data['model_id'] + target_data = DataFrame(input_data['target_data']) + metrics = input_data['metrics'] + interval_minutes = input_data['interval_minutes'] + + data_size = target_data.shape[0] + + output_data = [] + + diff = target_data['target'] - target_data['prediction'] + diff_squared = diff**2 + + self.info(f'Calculating simple metrics for model {model_id}: {metrics}', metadata) + + for metric in metrics: + if metric == 'rmse': + output_data.append({'metric': 'rmse', 'value': np.sqrt(np.mean(diff_squared))}) + elif metric == 'mse': + output_data.append({'metric': 'mse', 'value': np.mean(diff_squared)}) + elif metric == 'mae': + output_data.append({'metric': 'mae', 'value': np.mean(np.abs(diff))}) + elif metric == 'r2': + y_true = target_data['target'] + y_mean = np.mean(y_true) + + ss_res = np.sum(diff_squared) + ss_tot = np.sum((y_true - y_mean) ** 2) + + # Evita divisão por zero + if ss_tot == 0: + r2_score = 0.0 + else: + r2_score = 1 - (ss_res / ss_tot) + + output_data.append({'metric': 'r2', 'value': r2_score}) + + data = DataFrame(output_data) + data['model_id'] = model_id + data['timestamp'] = target_data['timestamp'].max() + data['data_size'] = data_size + data['interval_minutes'] = interval_minutes + + self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata) + + return data.to_dict(orient='records') diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py new file mode 100644 index 0000000..04ed18e --- /dev/null +++ b/laborious/activities/opc.py @@ -0,0 +1,515 @@ +from temporalio import activity, workflow + +with workflow.unsafe.imports_passed_through(): + import traceback + from collections.abc import Hashable + from typing import Any + + from pandas import DataFrame + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + from laborious.utils.repository.opc_repository import OpcRepository + +OPC_WRITTING_ERROR_CONFIDENCE = 12 +OPC_SESSION_BAD_CONFIDENCE = 14 +OPC_SESSION_BAD_COMMENT_PREFIX = 'OPC UA session/channel error:' +OPC_WRITTING_ERROR_MESSAGE = 'Some data could not be written to OPC servers' +OPC_RECONNECT_IN_PROGRESS_COMMENT = 'OPC UA reconnect in progress' +OPC_COMMENT_SEPARATOR = ' | ' + + +def _opc_session_bad_comment(opc_status: str | None) -> str: + status = opc_status or 'Unknown' + return f'{OPC_SESSION_BAD_COMMENT_PREFIX} {status}' + + +def _apply_opc_write_error( + error_info: dict[str, Any] | None, + session_bad_seen: bool, + session_bad_status: str | None, + reconnect_in_progress_seen: bool, +) -> tuple[bool, str | None, bool]: + """ + Update session/reconnect flags from an OPC write error payload. + + Args: + error_info: Repository error details, or None when the write succeeded. + session_bad_seen: Whether a session_bad error was seen so far. + session_bad_status: Last known OPC status for session errors. + reconnect_in_progress_seen: Whether reconnect_in_progress was seen so far. + + Return: + Updated (session_bad_seen, session_bad_status, reconnect_in_progress_seen). + """ + if not error_info: + return session_bad_seen, session_bad_status, reconnect_in_progress_seen + + kind = error_info.get('opc_error_kind') + if kind == 'session_bad': + return True, error_info.get('opc_status', session_bad_status), reconnect_in_progress_seen + if kind == 'reconnect_in_progress': + return session_bad_seen, session_bad_status, True + return session_bad_seen, session_bad_status, reconnect_in_progress_seen + + +class OPC(SientiaMonitoring): + """ + OPC server integration activities for real-time data export. + + This class provides comprehensive OPC UA client functionality for connecting + to multiple OPC servers and writing prediction data in real-time. It implements + secure communication with certificate-based authentication and automatic + reconnection capabilities. + + The class supports multiple OPC servers with individual configurations and + provides robust error handling and monitoring for production environments. + + Attributes: + opc_servers (dict): Configuration for multiple OPC servers + opc_repository (dict): Active OPC repository connections + logger (Logger): Logging and observability instance + notification_handler (NotificationHandler): Notification management instance + """ + + def __init__( + self, + opc_servers: dict[str, dict[str, Any]], + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + ): + self.logger = logger + self.notification_handler = notification_handler + self.opc_servers = opc_servers + + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + self.opc_repository: dict[str, OpcRepository] = {} + + async def init_opc(self): + """ + Initialize OPC server connections and establish communication channels. + + This method iterates through all configured OPC servers and attempts to + establish secure connections using certificate-based authentication. + Each server connection is managed independently, and connection failures + are reported through the notification system. + + The method performs the following operations: + 1. Creates OpcRepository instances for each configured server + 2. Establishes secure connections with certificate validation + 3. Reports connection success/failure through notifications + 4. Logs connection status for operational visibility + + Raises: + Exception: If OPC repository initialization fails or connection + establishment encounters critical errors + + Note: + Connection failures are logged and reported but do not prevent + the initialization of other OPC servers. Each server is handled + independently to ensure maximum availability. + """ + self.logger.info('Initializing OPC servers...') + for opc_id, server in self.opc_servers.items(): + self.opc_repository[opc_id] = OpcRepository( + opc_id=server['id'], + server_name=server['server_name'], + url=server['url'], + logger=self.logger, + server_uri=server['server_uri'], + cert_path=server['cert_path'], + private_key_path=server['private_key_path'], + server_cert_path=server['server_cert_path'], + notification_handler=self.notification_handler, + reconnection_interval=server['reconnection_interval'], + metrics_controller=self.metrics_controller, + ) + is_connected, error_data = await self.opc_repository[opc_id].connect() + if not is_connected: + await self.send_notification_async( + metadata={ + 'model_id': '-', + 'model_name': '-', + 'workflow_name': '-', + 'schedule_name': 'INITIALIZATION', + }, + notification_id=error_data['notification_id'], + message=error_data['message'], + block=error_data['block'], + level=error_data.get('level', NotificationLevel.ERROR), + attachment_content=error_data.get('attachment_content', None), + ) + else: + self.logger.info( + f'OPC server {opc_id}:{server["server_name"]} connected successfully.' + ) + + async def write_data( + self, + server_id: str, + tag: str, + data: Any, + data_type: str, + tag_type: str, + metadata: dict[str, Any], + ) -> tuple[float | None, dict[str, Any] | None]: + """ + Write data to a specific OPC server tag with comprehensive error handling. + + Return: + tuple[float | None, dict[str, Any] | None]: Response time on success, or + (None, error info_data) on repository failure. + """ + + try: + is_success, info_data = await self.opc_repository[server_id].write_data( + tag, data, data_type, metadata + ) + if not is_success: + await self.send_notification_async( + metadata=metadata, + notification_id=info_data['notification_id'], + message=info_data['message'], + block=info_data['block'], + level=info_data.get('level', NotificationLevel.ERROR), + attachment_content=info_data.get('attachment_content', None), + ) + return None, info_data + return info_data['response_time'], None + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR', + message=f'Error writing data to OPC server: {e}', + block='write_opc_data', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + raise e + + async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool: + """ + Validate that an OPC server is available and configured for write operations. + + This method checks if the specified OPC server exists in the active + repository and is available for data writing operations. It provides + immediate feedback for server availability and logs validation failures + for operational monitoring. + + Args: + server_id (str): Unique identifier for the OPC server to validate + metadata (dict[str, Any]): Context metadata for logging and notifications + + Returns: + bool: True if server is available, False otherwise + + Note: + Server validation failures are automatically reported through the + notification system with detailed information about available servers. + This helps operators quickly identify configuration issues. + """ + if self.opc_repository.get(server_id) is None: + message = f'OPC server {server_id} not found to perform write operation.' + await self.send_notification_async( + metadata=metadata, + notification_id='OPC_SERVER_NOT_FOUND', + message=message, + block='write_opc_data', + level=NotificationLevel.ERROR, + attachment_content=f'OPC servers: {list(self.opc_repository.keys())}', + ) + return False + return True + + async def _write_tags_from_config( + self, + server_id: str, + tags_config: dict[str, dict[str, Any]], + data: DataFrame, + data_column: str, + tag_type: str, + log_label: str, + metadata: dict[str, Any], + ) -> tuple[dict[str, float | None], bool, str | None, bool]: + """ + Write a group of OPC tags and collect response times and error flags. + + Args: + server_id: Target OPC server identifier. + tags_config: Tag name to configuration mapping. + data: DataFrame with prediction/confidence columns. + data_column: Column name whose first row value is written. + tag_type: Tag category passed to write_data ('prediction' or 'confidence'). + log_label: Human-readable label for success logs. + metadata: Context metadata for logging and notifications. + + Return: + (response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen) + """ + response_times: dict[str, float | None] = {} + session_bad_seen = False + session_bad_status: str | None = None + reconnect_in_progress_seen = False + + for tag, tag_config in tags_config.items(): + response_time, error_info = await self.write_data( + server_id=server_id, + tag=tag, + data=data.head(1)[data_column].values[0], + data_type=tag_config['data_type'], + tag_type=tag_type, + metadata=metadata, + ) + session_bad_seen, session_bad_status, reconnect_in_progress_seen = ( + _apply_opc_write_error( + error_info, + session_bad_seen, + session_bad_status, + reconnect_in_progress_seen, + ) + ) + if response_time is not None: + self.info( + f'{log_label} written to OPC server {server_id} for tag {tag}.', + metadata, + ) + response_times[tag] = response_time + + return response_times, session_bad_seen, session_bad_status, reconnect_in_progress_seen + + async def manage_output_tags( + self, + server_id: str, + config: dict[str, Any], + data: DataFrame, + metadata: dict[str, Any], + ) -> tuple[bool, dict[str, float | None], bool, str | None, bool]: + """ + Manage the writing of prediction and confidence data to OPC server tags. + + This method orchestrates the writing of multiple data types to OPC servers + based on configuration. It handles both prediction data and confidence + values independently, allowing for flexible tag configuration and + comprehensive error handling. + + The method supports two main tag types: + 1. Prediction tags: Write actual prediction values to configured OPC tags + 2. Confidence tags: Write confidence scores to separate OPC tags + + Args: + server_id (str): Unique identifier for the target OPC server + config (dict[str, Any]): OPC tag configuration containing: + - prediction_tags (dict, optional): Prediction tag configurations + - confidence_tags (dict, optional): Confidence tag configurations + data (DataFrame): DataFrame containing prediction and confidence data + metadata (dict[str, Any]): Context metadata for logging and notifications + success (bool): Current success status to maintain across operations + + Returns: + tuple[bool, int]: (overall_success, total_tags_written) + - overall_success: True if all configured tags were written successfully + - total_tags_written: Count of successfully written tags + """ + response_times: dict[str, float | None] = {} + session_bad_seen = False + session_bad_status: str | None = None + reconnect_in_progress_seen = False + + tag_groups = ( + ('prediction_tags', 'prediction', 'prediction', 'Prediction data'), + ('confidence_tags', 'prediction_confidence', 'confidence', 'Confidence data'), + ) + for config_key, data_column, tag_type, log_label in tag_groups: + if config_key not in config: + continue + ( + group_times, + group_session_bad, + group_status, + group_reconnect, + ) = await self._write_tags_from_config( + server_id=server_id, + tags_config=config[config_key], + data=data, + data_column=data_column, + tag_type=tag_type, + log_label=log_label, + metadata=metadata, + ) + response_times.update(group_times) + if group_session_bad: + session_bad_seen = True + session_bad_status = group_status or session_bad_status + if group_reconnect: + reconnect_in_progress_seen = True + + success = None not in response_times.values() + return ( + success, + response_times, + session_bad_seen, + session_bad_status, + reconnect_in_progress_seen, + ) + + @activity.defn(name='write_opc_data') + async def write_opc_data( + self, input_data: dict[str, Any] + ) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]: + """ + Write prediction and confidence data to OPC servers. The two writing + operations are optional and independent of each other. + + Args: + - input_data(dict[str, Any]): The input data. Contains the following keys: + - data(dict[str, Any]): The dataframe that contains the data to write + to the OPC servers. + - opc_output_config(dict[str, Any]): The OPC writing configuration. + The keys are the OPC server names and the values contain: + - prediction_tags(dict[str, Any]): The tags to write to the OPC servers. + - confidence_tags(dict[str, Any]): The tags to write to the OPC servers. + + Returns: + - dict[Any, Any]: The data that was written to the OPC servers. + + """ + metadata = input_data['metadata'] + self.info('Writing data to OPC servers...', metadata) + data = DataFrame(input_data['data']) + opc_output_config = input_data['opc_output_config'] + self.info(f'Data to write: {data.size} rows', metadata) + + success = True + session_bad_seen = False + session_bad_status: str | None = None + reconnect_in_progress_seen = False + + metrics: dict[str, dict[str, float | None]] = {} + + for server_id, config in opc_output_config.items(): + if not await self.validate_server(server_id, metadata): + success = False + continue + + ( + local_success, + local_response_times, + local_session_bad, + local_status, + local_reconnect_in_progress, + ) = await self.manage_output_tags(server_id, config, data, metadata) + metrics[server_id] = local_response_times + local_count = len(local_response_times) + success = success and local_success + if local_session_bad: + session_bad_seen = True + session_bad_status = local_status or session_bad_status + if local_reconnect_in_progress: + reconnect_in_progress_seen = True + + self.info( + f'Process completed for OPC server {server_id}: {local_count} of {len(config.get("prediction_tags", []))} prediction tags and {len(config.get("confidence_tags", []))} confidence tags', + metadata, + ) + + return ( + self.process_confidence( + data, + success, + metadata, + session_bad=session_bad_seen, + opc_status=session_bad_status, + reconnect_in_progress=reconnect_in_progress_seen, + ), + metrics, + ) + + def process_confidence( + self, + data: DataFrame, + success: bool, + metadata: dict[str, Any], + *, + session_bad: bool = False, + opc_status: str | None = None, + reconnect_in_progress: bool = False, + ) -> dict[Hashable, Any]: + """ + Process prediction confidence based on OPC write operation success. + + This method updates the prediction confidence values in the DataFrame + based on the success status of OPC server write operations. If any + write operations failed, it sets the confidence to a predefined error + value to indicate data quality issues. + + The method implements a confidence degradation strategy: + - Success: Maintains original confidence values + - Failure: Sets confidence to error value for operational awareness + + Args: + data (DataFrame): DataFrame containing prediction and confidence data + success (bool): Overall success status of OPC write operations + metadata (dict[str, Any]): Context metadata for logging and notifications + + Returns: + dict[Any, Any]: Processed data as a dictionary with updated confidence values + + Note: + The error confidence value (OPC_WRITTING_ERROR_CONFIDENCE = 12) is + used to indicate that data was not successfully exported to OPC servers. + This allows downstream systems to handle data quality appropriately. + """ + + if not success: + comment_parts: list[str] = [] + confidence = OPC_WRITTING_ERROR_CONFIDENCE + + if session_bad: + comment_parts.append(_opc_session_bad_comment(opc_status)) + confidence = OPC_SESSION_BAD_CONFIDENCE + if reconnect_in_progress: + comment_parts.append(OPC_RECONNECT_IN_PROGRESS_COMMENT) + confidence = OPC_SESSION_BAD_CONFIDENCE + if not comment_parts: + comment_parts.append(OPC_WRITTING_ERROR_MESSAGE) + + comments = OPC_COMMENT_SEPARATOR.join(comment_parts) + data['prediction_confidence'] = confidence + data['comments'] = comments + self.debug( + f'OPC write issues, confidence={confidence}, comments={comments}', + metadata, + ) + else: + self.debug('Data written to OPC servers successfully.', metadata) + + return data.to_dict() + + async def aclose(self): + """ + Gracefully shutdown all OPC server connections and cleanup resources. + + This method ensures proper cleanup of all active OPC server connections + by calling the disconnect method on each repository instance. It's + designed to be called during application shutdown to prevent resource + leaks and ensure clean termination. + + The method performs the following cleanup operations: + 1. Iterates through all active OPC repository connections + 2. Calls disconnect() on each repository instance + 3. Allows for graceful connection termination + 4. Prevents resource leaks and connection hanging + + Note: + This method should be called during application shutdown to ensure + proper cleanup. It handles all active connections regardless of + their current state and provides a clean shutdown experience. + """ + for opc in self.opc_repository.values(): + await opc.disconnect() diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py new file mode 100644 index 0000000..7c67466 --- /dev/null +++ b/laborious/activities/storage.py @@ -0,0 +1,210 @@ +from temporalio import activity, workflow + +from laborious.utils.repository.minio_manager import MinioManager + +with workflow.unsafe.imports_passed_through(): + # Extend the Temporal Postgres activities for convenient query -> MinIO export + import traceback + from datetime import timedelta + from typing import Any + + import pandas as pd + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.notifications.models import NotificationLevel + from sientia_do.observability.logger import Logger + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.repository.minio_repository import MinioRepository + from sientia_do.temporal.activities.postgres import Postgres + from sientia_do.temporal.constants import now + + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + +_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'}) + + +class Storage(Postgres, MinioManager): + """ + Extensions for Postgres activities with a helper to export query results + directly to MinIO as Parquet and return the object name. + """ + + minio_repository: MinioRepository | None = None + + def __init__( + self, + host: str, + port: int, + user: str, + password: str, + dbname: str, + min_connections: int, + max_connections: int, + retention_hours: int = 24, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + self.retention_hours = retention_hours + Postgres.__init__( + self, + host=host, + port=port, + user=user, + password=password, + dbname=dbname, + min_connections=min_connections, + max_connections=max_connections, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) + + @activity.defn(name='load_query_with_minio_offload') + async def load_query_with_minio_offload( + self, input_data: dict[str, Any] + ) -> MinioDataFramePayload: + """ + Run the custom SQL load, then return a MinIO-aware dataframe wire dict. + + Args (input_data): + metadata (dict): Workflow metadata (same as load_custom_query). + query (str): SQL query. + datetime_columns (list[str], optional): Datetime column names. + model_name (str): Model name for object key basename. + key_prefix (str, optional): Directory prefix inside the bucket. + size_threshold_bytes (int, optional): Override env offload threshold. + + Returns: + dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure. + """ + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + + metadata: dict = input_data.get('metadata', {}) + model_name = input_data['model_name'] + + rows = await self.load_custom_query( + input_data, + ) + if not rows: + self.error( + 'load_query_with_minio_offload failed: No data returned from query', metadata + ) + dataframe = None + else: + dataframe = pd.DataFrame(rows) + + return await MinioDataFramePayload.from_dataframe( + dataframe, + minio_repo=self.minio_repository, + workflow_metadata=metadata, + model_name=model_name, + operation='initial', + logger=self.logger, + ) + + @activity.defn(name='export_payload_to_postgres') + async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict: + """ + Export a payload to PostgreSQL. + """ + metadata = input_data.get('metadata') + payload = MinioDataFramePayload.from_dict(input_data['data']) + data = await payload.retrieve(self.minio_repository, metadata) + + return await self.export_data_to_postgres( + { + **input_data, + 'data': data, + } + ) + + @activity.defn(name='cleanup_minio_objects_expired') + async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]: + """ + Delete objects under the given prefixes that are older than the retention window. + + Args (input_data): + metadata (dict): Workflow metadata for logging and metrics. + prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix). + + Returns: + dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``. + """ + if self.minio_repository is None: + raise ValueError('Minio repository not initialized') + + metadata = input_data.get('metadata', {}) + payload = MinioDataFramePayload.from_dict(input_data['data']) + prefix = payload.cleanup_prefix() + base = now() + cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta( + hours=self.retention_hours + ) + + report: dict[str, Any] = { + 'failed': {}, + 'deleted': {}, + 'failed_count': 0, + 'deleted_count': 0, + } + try: + keys = await self.minio_repository.list_objects( + prefix=prefix, + recursive=True, + metadata=metadata, + ) + for key in keys: + try: + ts = MinioDataFramePayload.parse_object_timestamp(key) + if ts is None: + continue + if ts >= cutoff: + continue + await self.minio_repository.delete_file( + object_name=key, + metadata=metadata, + ) + except Exception as e: + report['failed'][key] = { + 'success': False, + 'message': str(e), + } + report['failed_count'] += 1 + continue + report['deleted'][key] = { + 'success': True, + 'message': 'Deleted', + } + report['deleted_count'] += 1 + except Exception as e: + trace = traceback.format_exc() + await self.send_notification_async( + metadata=metadata, + notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED', + message=f'Error cleaning up MinIO objects: {e}', + block='cleanup_minio_objects_expired', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + self.error(trace, metadata) + else: + # Cleanup success is expected in normal flow; avoid noisy INFO notifications + # that do not impact behavior and can flood observability in test runs. + self.info('MinIO objects cleaned up successfully', metadata) + + return report + + def close(self) -> None: + """Close Storage resources (MinIO client and Postgres engine).""" + if hasattr(self, 'engine'): + Postgres.close(self) + MinioManager.close(self) + + def __del__(self): + self.close() diff --git a/laborious/metrics.py b/laborious/metrics.py new file mode 100644 index 0000000..3f4de2c --- /dev/null +++ b/laborious/metrics.py @@ -0,0 +1,239 @@ +""" +Laborious Metrics Module + +This module defines all Prometheus metrics used by the Sientia DataOps Laborious system +for monitoring and observability. The metrics provide insights into system performance, +prediction quality, and operational health. + +The metrics are designed to be scraped by Prometheus and can be visualized in +Grafana or other monitoring dashboards to provide real-time visibility into +the system's operation. + +Key Metric Categories: +- Application Health: Overall system status and availability +- Prediction Operations: Count and performance of prediction operations +- Data Quality: Confidence levels and validation results +- Export Operations: Database and OPC export performance +- Response Times: Performance monitoring for various operations + +Metric Labels: +- pod_id: Kubernetes pod identifier for multi-instance deployments +- runtime: Runtime / environment identifier (matches ``RUNTIME`` env, see ``SientiaMonitoring``) +- model_name: Name of the ML model being used +- workflow_name: Name of the prediction pipeline +- opc_server_id: Identifier for OPC server operations +""" + +from prometheus_client import Counter, Gauge, Histogram +from sientia_do.observability.metrics import CORE_LABELS + +# Application health metric +APP_UP = Gauge( + 'app_up', + 'Indicates if the application is running (1) or shutting down (0)', + ['pod_id'], +) + +# Prediction operation metrics +PREDICTIONS_WRITTEN_COUNT = Counter( + 'laborious_predictions_written_count', + 'Number of predictions written to the database table predictions', + CORE_LABELS, +) + +# Prediction quality metrics +PREDICTION_CONFIDENCE_MONITOR = Gauge( + 'laborious_prediction_confidence_monitor', + 'Current confidence of each prediction', + CORE_LABELS, +) + +# Prediction total response time +PREDICTION_RESPONSE_TIME_MONITOR = Histogram( + 'laborious_prediction_response_time_monitor', + 'Current response time of each prediction', + CORE_LABELS, + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + + +# ================== OPC metrics ================== + + +PREDICTION_OPC_WRITING_COUNT = Counter( + 'laborious_prediction_opc_writing_count', + 'Number of predictions written to the OPC server', + [*CORE_LABELS, 'opc_server_id', 'tag'], +) + +PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram( + 'laborious_prediction_opc_writing_response_time_monitor', + 'Current response time of each prediction written to the OPC server', + [*CORE_LABELS, 'opc_server_id', 'tag'], + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + +OPC_CONNECTIONS_TOTAL = Counter( + 'opc_connections_initiated_total', + 'Total connection attempts to OPC servers', + ['pod_id', 'server_name'], +) +OPC_CONNECTIONS_FAILED = Counter( + 'opc_connections_failed_total', + 'Total failed connection attempts to OPC servers', + ['pod_id', 'server_name'], +) +OPC_CONNECTION_STATUS = Gauge( + 'opc_connection_status', + 'Connection status with the OPC server (1=connected, 0=disconnected)', + ['pod_id', 'server_name', 'server_url'], +) + +_OPC_SESSION_DEBUG_LABELS = ['pod_id', 'server_name', 'runtime', 'opc_server_id', 'session_id'] + +OPC_SESSION_CREATED_TOTAL = Counter( + 'opc_session_created_total', + 'OPC UA sessions established (after successful connect)', + _OPC_SESSION_DEBUG_LABELS, +) + +OPC_SESSION_CLOSED_TOTAL = Counter( + 'opc_session_closed_total', + 'OPC UA client disconnects completed (session tear-down initiated)', + _OPC_SESSION_DEBUG_LABELS, +) + +OPC_SESSION_REVISED_TIMEOUT_MS = Gauge( + 'opc_session_revised_timeout_milliseconds', + 'Server-revised OPC UA session timeout (RevisedSessionTimeout) in ms after connect', + _OPC_SESSION_DEBUG_LABELS, +) + +OPC_WRITE_ATTEMPT_LABELS = [*_OPC_SESSION_DEBUG_LABELS, 'model_id', 'model_name', 'result'] + +OPC_WRITE_ATTEMPTS_TOTAL = Counter( + 'opc_write_attempts_total', + 'OPC UA write attempts with session and outcome (result=OK or exception class name)', + OPC_WRITE_ATTEMPT_LABELS, +) + +OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL = Counter( + 'opc_write_inter_arrival_over_session_timeout_total', + 'Successful writes where seconds since the previous successful write exceeded RevisedSessionTimeout (ms)', + _OPC_SESSION_DEBUG_LABELS, +) + +# ================== Model metrics ================== + +MODEL_READ_LAG = Histogram( + 'laborious_model_read_lag', + 'Lag between the start and read of read operations', + CORE_LABELS, + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + +MODEL_WRITE_LAG = Histogram( + 'laborious_model_write_lag', + 'Lag between the start and end of write operations', + CORE_LABELS, + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + +MODEL_READ_COUNT = Counter( + 'laborious_model_read_count', + 'Number of reads from the model', + CORE_LABELS, +) + +MODEL_WRITE_COUNT = Counter( + 'laborious_model_write_count', + 'Number of writes to the model', + CORE_LABELS, +) + +MODEL_READ_ERROR_COUNT = Counter( + 'laborious_model_read_error_count', + 'Number of errors reading from the model', + CORE_LABELS, +) + +MODEL_WRITE_ERROR_COUNT = Counter( + 'laborious_model_write_error_count', + 'Number of errors writing to the model', + CORE_LABELS, +) + +MODEL_ANALYZE_LAG = Histogram( + 'laborious_model_analyze_lag', + 'Lag between the start and end of analyze operations', + CORE_LABELS, + buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], +) + +MODEL_ANALYZE_COUNT = Counter( + 'laborious_model_analyze_count', + 'Number of analyze operations', + CORE_LABELS, +) + +MODEL_ANALYZE_ERROR_COUNT = Counter( + 'laborious_model_analyze_error_count', + 'Number of errors during analyze operations', + CORE_LABELS, +) + + +# ================== Model import metrics ================== +# +# The failing step and the error code travel in `operation_type`, the one label the shared +# `CORE_LABELS` leaves free. Nothing else about a failure is a label: an error reason is a sentence +# and an exception is a log line, neither of which belongs in a metric's cardinality. + +MODEL_IMPORT_STARTED_COUNT = Counter( + 'laborious_model_import_started_count', + 'Number of model imports that claimed their import log row', + CORE_LABELS, +) + +MODEL_IMPORT_COMPLETED_COUNT = Counter( + 'laborious_model_import_completed_count', + 'Number of model imports that reached a registered version and a model document', + CORE_LABELS, +) + +MODEL_IMPORT_ERROR_COUNT = Counter( + 'laborious_model_import_error_count', + 'Number of failed model imports, with the failing step or error code in operation_type', + CORE_LABELS, +) + +MODEL_IMPORT_KDF_DURATION = Histogram( + 'laborious_model_import_kdf_duration_seconds', + 'Time spent deriving the bundle key with Argon2id', + CORE_LABELS, + buckets=[0.5, 1.0, 2.0, 3.0, 5.0, 8.0, 13.0, 21.0], +) + +MODEL_IMPORT_STEP_DURATION = Histogram( + 'laborious_model_import_step_duration_seconds', + 'Duration of each model import step, named in operation_type', + CORE_LABELS, + buckets=[0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 300.0], +) + + +# ================== PI Web API metrics ================== + +PI_WEB_API_LABELS = [*CORE_LABELS, 'tag_name'] + +PI_WEB_API_PREDICTION_WRITTEN_COUNT = Counter( + 'laborious_pi_web_api_prediction_written_count', + 'Number of predictions written to the PI Web API', + PI_WEB_API_LABELS, +) + +PI_WEB_API_PREDICTION_WRITTEN_ERROR_COUNT = Counter( + 'laborious_pi_web_api_prediction_written_error_count', + 'Number of errors writing predictions to the PI Web API', + PI_WEB_API_LABELS, +) diff --git a/laborious/utils/__init__.py b/laborious/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/bundle/__init__.py b/laborious/utils/bundle/__init__.py new file mode 100644 index 0000000..77e7945 --- /dev/null +++ b/laborious/utils/bundle/__init__.py @@ -0,0 +1,27 @@ +"""Reading `.sientia` bundles: the wire format, the seven gates and the step vocabulary. + +There is no writer here. The producer is the Streamlit app in `sientia-projects-templates` +(`app/src/operations/model_export/`), and a second production writer for a format with one producer +would be a fork waiting to happen. What stands in for it in the tests is +`tests/helpers/bundle_factory.py` plus the golden fixtures under `tests/fixtures/bundle/`. +""" + +from laborious.utils.bundle.format import BundleFormatError, BundleHeader +from laborious.utils.bundle.reader import ( + BundleGateError, + BundleLimits, + BundleReader, + OpenedBundle, +) +from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep + +__all__ = [ + 'GATE_TO_STEP', + 'BundleFormatError', + 'BundleGateError', + 'BundleHeader', + 'BundleLimits', + 'BundleReader', + 'ImportStep', + 'OpenedBundle', +] diff --git a/laborious/utils/bundle/format.py b/laborious/utils/bundle/format.py new file mode 100644 index 0000000..f07d2f8 --- /dev/null +++ b/laborious/utils/bundle/format.py @@ -0,0 +1,216 @@ +"""The `.sientia` wire format, as the producer writes it. + +Every constant here is copied from the producer that emits the bundles this reader opens — the +Streamlit app in `sientia-projects-templates`, modules +`app/src/operations/model_export/crypto.py` and `app/src/operations/model_export/bundle.py`. The +producer file each value came from is named next to it. The golden fixtures under +`tests/fixtures/bundle/` are the anti-drift device: if the producer's layout moves, the tests here +fail against the committed bytes instead of failing in production. + +This module deliberately imports nothing from Laborious, Temporal, MLflow or libsodium: it is pure +`struct` and `hashlib` over bytes, so it stays liftable into a shared package if the format ever +gets one (the cancelled QTZPOC-15). `reader.py` is where the crypto lives. +""" + +from __future__ import annotations + +import hashlib +import struct +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +# --- Producer: app/src/operations/model_export/crypto.py --- +MAGIC: Final[bytes] = b'SIENTIA1' +FORMAT_VERSION: Final[int] = 1 +AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM: Final[int] = 1 +KDF_ID_ARGON2ID: Final[int] = 1 +SALT_SIZE: Final[int] = 16 +DIGEST_SIZE: Final[int] = 32 +RESERVED_SIGNATURE_SIZE: Final[int] = 64 +# `KDF_MEMLIMIT_BYTES` / `KDF_OPSLIMIT` on the producer: what it writes is also the ceiling the +# reader refuses to exceed, so header bytes cannot buy an arbitrary allocation on the worker. +MAX_KDF_MEMLIMIT: Final[int] = 256 * 1024 * 1024 # 268435456 +MAX_KDF_OPSLIMIT: Final[int] = 3 +# `KDF_PARALLELISM_RESERVED`: libsodium's `crypto_pwhash` has no lane parameter, so the field is +# reserved and must be exactly 1. +PARALLELISM: Final[int] = 1 +CHUNK_SIZE: Final[int] = 64 * 1024 # 65536, fixed by the format +KEY_SIZE: Final[int] = 32 # crypto_secretstream_xchacha20poly1305_KEYBYTES + +# magic(8s) version(B) aead_id(B) kdf_id(B) salt(16s) memlimit(Q) opslimit(Q) +# parallelism(B) digest(32s) reserved_signature(64s) — producer's `_HEADER_STRUCT`. +HEADER_STRUCT: Final[struct.Struct] = struct.Struct('>8sBBB16sQQB32s64s') +HEADER_SIZE: Final[int] = HEADER_STRUCT.size # 140 + +# --- libsodium constants the framing depends on --- +# Not producer constants, but constants of the primitive the producer uses. Kept as literals so +# this module stays import-free; `tests/laborious/utils/bundle/test_format.py` asserts they equal +# the values `nacl.bindings` reports, which is what makes the literals safe. +STREAM_HEADER_SIZE: Final[int] = 24 # crypto_secretstream_xchacha20poly1305_HEADERBYTES +ABYTES: Final[int] = 17 # crypto_secretstream_xchacha20poly1305_ABYTES + +# The pre-upload check prefix of `08-security-and-encryption.md` § 5.3: enough bytes to hold the +# header, the stream header and one full chunk. Derived from constants, never read from the file. +PREFIX_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + CHUNK_SIZE + ABYTES # 65717 + +# The digest covers the chunk stream only. The producer computes `sha256(ciphertext)` over the +# joined chunks and *then* writes `header + stream_header + ciphertext`, so on disk the covered +# range starts after both headers. Verified against a bundle produced by the producer's own code: +# `sha256(file[164:])` matches the header field and `sha256(file[140:])` does not. +DIGEST_OFFSET: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE # 164 + +MIN_FILE_SIZE: Final[int] = HEADER_SIZE + STREAM_HEADER_SIZE + +SUPPORTED_FORMAT_VERSIONS: Final[frozenset[int]] = frozenset({FORMAT_VERSION}) +SUPPORTED_AEAD_IDS: Final[frozenset[int]] = frozenset({AEAD_ID_XCHACHA20POLY1305_SECRETSTREAM}) +SUPPORTED_KDF_IDS: Final[frozenset[int]] = frozenset({KDF_ID_ARGON2ID}) + +_HASH_BLOCK_SIZE: Final[int] = 1024 * 1024 + + +class BundleFormatError(Exception): + """The bytes are not a `.sientia` bundle this reader implements. + + Raised for every header-level rejection: wrong magic, unknown version, unknown primitive ids, + cost parameters above the ceiling, a reserved field carrying a value, a file too short to hold + a header, or a chunk stream whose digest does not match the header's. `reader.py` turns this + into a gate-2 rejection; nothing here knows about gates or import steps. + """ + + +@dataclass(frozen=True) +class BundleHeader: + """The 140-byte plaintext header, parsed and validated.""" + + format_version: int + aead_id: int + kdf_id: int + salt: bytes + kdf_memlimit_bytes: int + kdf_opslimit: int + kdf_parallelism: int + ciphertext_digest: bytes + reserved_signature: bytes + + +def parse_header(raw: bytes) -> BundleHeader: + """Parse and validate the plaintext header. + + Everything the reader trusts about the file's shape is decided here, before any offset past + the header is used and before a key is derived. + + Args: + raw: at least `HEADER_SIZE` bytes read from the start of the object. + + Returns: + BundleHeader: the validated header. + + Raises: + BundleFormatError: on any invalid or unsupported field. + """ + if len(raw) < HEADER_SIZE: + raise BundleFormatError( + f'file is shorter than the {HEADER_SIZE}-byte header ({len(raw)} bytes available)' + ) + + ( + magic, + format_version, + aead_id, + kdf_id, + salt, + kdf_memlimit_bytes, + kdf_opslimit, + kdf_parallelism, + ciphertext_digest, + reserved_signature, + ) = HEADER_STRUCT.unpack(raw[:HEADER_SIZE]) + + if magic != MAGIC: + raise BundleFormatError('file does not start with the expected bundle marker') + if format_version not in SUPPORTED_FORMAT_VERSIONS: + raise BundleFormatError(f'unsupported bundle format version: {format_version}') + if aead_id not in SUPPORTED_AEAD_IDS: + raise BundleFormatError(f'unsupported encryption identifier: {aead_id}') + if kdf_id not in SUPPORTED_KDF_IDS: + raise BundleFormatError(f'unsupported key derivation identifier: {kdf_id}') + if kdf_memlimit_bytes > MAX_KDF_MEMLIMIT: + raise BundleFormatError( + f'declared key derivation memory {kdf_memlimit_bytes} exceeds the ' + f'{MAX_KDF_MEMLIMIT} ceiling' + ) + if kdf_opslimit > MAX_KDF_OPSLIMIT: + raise BundleFormatError( + f'declared key derivation passes {kdf_opslimit} exceeds the {MAX_KDF_OPSLIMIT} ceiling' + ) + if kdf_memlimit_bytes <= 0 or kdf_opslimit <= 0: + raise BundleFormatError('key derivation parameters must be positive') + if kdf_parallelism != PARALLELISM: + raise BundleFormatError( + f'reserved parallelism field must be {PARALLELISM}, found {kdf_parallelism}' + ) + if reserved_signature != bytes(RESERVED_SIGNATURE_SIZE): + raise BundleFormatError('reserved signature field must be empty in this format version') + + return BundleHeader( + format_version=format_version, + aead_id=aead_id, + kdf_id=kdf_id, + salt=salt, + kdf_memlimit_bytes=kdf_memlimit_bytes, + kdf_opslimit=kdf_opslimit, + kdf_parallelism=kdf_parallelism, + ciphertext_digest=ciphertext_digest, + reserved_signature=reserved_signature, + ) + + +def read_header(path: Path | str) -> BundleHeader: + """Read and validate the header of the file at `path`. + + Raises: + BundleFormatError: when the file cannot hold a header plus a stream header, or when the + header itself is invalid. + """ + file_path = Path(path) + size = file_path.stat().st_size + if size < MIN_FILE_SIZE: + raise BundleFormatError( + f'file is {size} bytes, smaller than the {MIN_FILE_SIZE} bytes a bundle needs for its ' + 'header and encrypted stream header' + ) + with file_path.open('rb') as handle: + return parse_header(handle.read(HEADER_SIZE)) + + +def compute_ciphertext_digest(path: Path | str) -> bytes: + """SHA-256 of the chunk stream — the bytes from `DIGEST_OFFSET` to the end of the file.""" + digest = hashlib.sha256() + with Path(path).open('rb') as handle: + handle.seek(DIGEST_OFFSET) + while block := handle.read(_HASH_BLOCK_SIZE): + digest.update(block) + return digest.digest() + + +def verify_ciphertext_digest(path: Path | str, header: BundleHeader) -> None: + """Compare the file's chunk-stream digest against the header's. + + Raises: + BundleFormatError: when they differ — a truncated, overwritten or corrupted transfer. This + runs before any key derivation, so a damaged file never costs an Argon2id pass. + """ + if compute_ciphertext_digest(path) != header.ciphertext_digest: + raise BundleFormatError( + 'the encrypted content does not match the digest recorded in the bundle header' + ) + + +def file_digest(path: Path | str) -> str: + """Hex SHA-256 of the whole file, for the `expected_digest` comparison of gate 1.""" + digest = hashlib.sha256() + with Path(path).open('rb') as handle: + while block := handle.read(_HASH_BLOCK_SIZE): + digest.update(block) + return digest.hexdigest() diff --git a/laborious/utils/bundle/reader.py b/laborious/utils/bundle/reader.py new file mode 100644 index 0000000..61fcd90 --- /dev/null +++ b/laborious/utils/bundle/reader.py @@ -0,0 +1,473 @@ +"""Reading a `.sientia` bundle: the seven gates, in order, in one function. + +`BundleReader.open()` is the only entry point. Gate ordering is a property of that function rather +than of the caller, so it cannot be reordered by accident and can be unit-tested without Temporal: + +1. object-level limits (`check_object_limits`, called by the download activity before the body is + fetched — the only gate that does not need the file) +2. header validation, including the chunk-stream digest, before any key derivation +3. AEAD decryption — the integrity check; a wrong password and a tampered file fail identically +4. archive inspection over the central directory, writing nothing to disk +5. extraction into a fresh, isolated directory +6. structure and schema of the extracted tree +7. content policy over the file set + +Nothing here logs, notifies or emits a metric: rejections are raised as `BundleGateError` and the +activity layer owns the one failure channel. +""" + +from __future__ import annotations + +import json +import os +import shutil +import stat +import time +import unicodedata +import zipfile +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import nacl.bindings as sodium +import nacl.exceptions + +from laborious.utils.bundle import format as fmt +from laborious.utils.bundle.steps import GATE_TO_STEP, ImportStep + +# Gate 6 — the layout `bundle.build_zip` writes and `model_repository.py` reads back. +METADATA_FILE_NAME = 'metadata.json' +ARTIFACTS_DIR_NAME = 'artifacts' +REQUIRED_TOP_LEVEL_KEYS = frozenset({'parameters', 'metadata'}) +REQUIRED_METADATA_FIELDS = ( + 'model_name', + 'experiment_name', + 'run_id', + 'model_version', + 'model_project', + 'export_timestamp', +) +REQUIRED_PARAMETERS = ('target_variable',) +REQUIRED_MODEL_DIRS = ('prediction_model', 'data_model') +REQUIRED_MODEL_FILES = ('MLmodel', 'model.pkl') + +# The import record's `experiment_name` column is `VARCHAR(50)` with a length-3 floor, and the name +# is written onto the record before provisioning starts — so the column's limits are enforced here, +# not three activities later as a constraint violation. +EXPERIMENT_NAME_MIN_LENGTH = 3 +EXPERIMENT_NAME_MAX_LENGTH = 50 + +# Gate 7 — an MLflow 2.x artifact tree for this platform: the model directories' known files, the +# model card, the CSVs the training template logs and the bundle's own metadata. +ALLOWED_FILE_NAMES = frozenset( + { + 'MLmodel', + 'metadata.json', + 'model.pkl', + 'conda.yaml', + 'python_env.yaml', + 'requirements.txt', + 'model_card.json', + 'model_card.svg', + } +) +ALLOWED_FILE_SUFFIXES = frozenset({'.csv', '.json', '.yaml', '.yml', '.txt', '.pkl', '.svg', '.md'}) + +_DECRYPT_FAILURE_MESSAGE = ( + 'could not open this bundle: wrong password, or the file is damaged or was altered' +) + + +class BundleGateError(Exception): + """A gate rejected the bundle. + + Carries the gate number and the import step that gate maps onto, so the workflow never has to + guess which step to record. + """ + + def __init__(self, message: str, *, gate: int, step: ImportStep | None = None) -> None: + super().__init__(message) + self.gate = gate + self.step = step if step is not None else GATE_TO_STEP[gate] + + +@dataclass(frozen=True) +class BundleLimits: + """Ceilings enforced by gates 1 and 4. Values come from `build_import_config()`.""" + + max_object_bytes: int = 1024 * 1024 * 1024 + max_entries: int = 5000 + max_uncompressed_bytes: int = 4 * 1024 * 1024 * 1024 + max_compression_ratio: float = 200.0 + object_prefix: str = 'imported_models/' + object_suffix: str = '.sientia' + + +@dataclass(frozen=True) +class OpenedBundle: + """What a bundle yields once every gate has passed. Nothing secret is carried.""" + + extracted_dir: Path + metadata: dict[str, Any] = field(default_factory=dict) + parameters: dict[str, Any] = field(default_factory=dict) + digest: str = '' + header: fmt.BundleHeader | None = None + # How long Argon2id took, so the worker's most expensive step is measurable without the + # activity layer having to reach inside the reader. + kdf_seconds: float = 0.0 + + +class BundleReader: + """Opens `.sientia` bundles produced by the exporter in `sientia-projects-templates`.""" + + def __init__(self, limits: BundleLimits | None = None) -> None: + self.limits = limits or BundleLimits() + + # ------------------------------------------------------------------ gate 1 + + def check_object_limits(self, *, size: int, key: str) -> None: + """Gate 1, object level: reject before the object's body is fetched. + + Args: + size: the object's size as reported by `stat_object`. + key: the object key. + + Raises: + BundleGateError: oversized object, key outside the import prefix, or wrong suffix. + """ + if size > self.limits.max_object_bytes: + raise BundleGateError( + f'uploaded object is {size} bytes, above the ' + f'{self.limits.max_object_bytes} byte limit', + gate=1, + ) + if size <= 0: + raise BundleGateError('uploaded object is empty', gate=1) + if not key.startswith(self.limits.object_prefix): + raise BundleGateError('uploaded object is not under the import prefix', gate=1) + if not key.endswith(self.limits.object_suffix): + raise BundleGateError( + f'uploaded object does not end in {self.limits.object_suffix}', gate=1 + ) + + # --------------------------------------------------------------------- open + + def open( + self, + encrypted_path: Path | str, + password: str, + extract_root: Path | str, + work_dir: Path | str | None = None, + ) -> OpenedBundle: + """Run gates 2 to 7 in order and return the opened bundle. + + Args: + encrypted_path: the downloaded `.sientia` file. + password: the plaintext bundle password. Never stored, never logged, never returned. + extract_root: the directory to extract into. It MUST NOT exist yet (gate 5). + work_dir: where the decrypted zip is written. Defaults to `extract_root`'s parent. + + Returns: + OpenedBundle: extracted tree, `metadata`, `parameters`, digest and header. + + Raises: + BundleGateError: with the gate that rejected the bundle and its import step. + """ + encrypted = Path(encrypted_path) + root = Path(extract_root) + staging = Path(work_dir) if work_dir is not None else root.parent + + header = self._gate_2_header(encrypted) + zip_path, kdf_seconds = self._gate_3_decrypt(encrypted, password, header, staging) + try: + self._gate_4_inspect(zip_path) + self._gate_5_extract(zip_path, root) + metadata, parameters = self._gate_6_structure(root) + self._gate_7_content_policy(root) + except BaseException: + # A rejection after decryption leaves nothing extracted behind (gate 5's contract) and + # never leaves a decrypted zip on the worker for the next import to trip over. + shutil.rmtree(root, ignore_errors=True) + zip_path.unlink(missing_ok=True) + raise + zip_path.unlink(missing_ok=True) + + return OpenedBundle( + extracted_dir=root, + metadata=metadata, + parameters=parameters, + digest=header.ciphertext_digest.hex(), + header=header, + kdf_seconds=kdf_seconds, + ) + + # ------------------------------------------------------------------ gate 2 + + def _gate_2_header(self, encrypted: Path) -> fmt.BundleHeader: + """Parse the header and check the chunk-stream digest, before deriving anything.""" + try: + header = fmt.read_header(encrypted) + fmt.verify_ciphertext_digest(encrypted, header) + except fmt.BundleFormatError as error: + raise BundleGateError(str(error), gate=2) from error + except OSError as error: + raise BundleGateError('downloaded bundle could not be read', gate=2) from error + return header + + # ------------------------------------------------------------------ gate 3 + + def _gate_3_decrypt( + self, encrypted: Path, password: str, header: fmt.BundleHeader, staging: Path + ) -> tuple[Path, float]: + """Derive the key with the header's parameters and decrypt the stream to a file. + + The AEAD *is* the integrity check: a wrong password and a tampered byte fail here with the + same message, which never claims to distinguish them. + """ + staging.mkdir(parents=True, exist_ok=True) + zip_path = staging / 'bundle.zip' + try: + started = time.monotonic() + key = self._derive_key(password, header) + kdf_seconds = time.monotonic() - started + self._pull_stream(encrypted, key, zip_path) + except BundleGateError: + zip_path.unlink(missing_ok=True) + raise + except (nacl.exceptions.CryptoError, ValueError, RuntimeError) as error: + zip_path.unlink(missing_ok=True) + raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) from error + except OSError as error: + zip_path.unlink(missing_ok=True) + raise BundleGateError('decrypted bundle could not be written', gate=3) from error + return zip_path, kdf_seconds + + @staticmethod + def _derive_key(password: str, header: fmt.BundleHeader) -> bytes: + """Argon2id over the NFC-normalised password, with the header's own cost parameters. + + The password is normalised and **not** trimmed, matching the producer's `_derive_key`: the + same password typed on another OS must derive the same key, while a trailing space is part + of the secret. + """ + normalized = unicodedata.normalize('NFC', password) + return sodium.crypto_pwhash_alg( + fmt.KEY_SIZE, + normalized.encode('utf-8'), + header.salt, + header.kdf_opslimit, + header.kdf_memlimit_bytes, + sodium.crypto_pwhash_ALG_ARGON2ID13, + ) + + @staticmethod + def _pull_stream(encrypted: Path, key: bytes, zip_path: Path) -> None: + """Decrypt the chunk stream into `zip_path`, requiring the FINAL tag.""" + chunk_size = fmt.CHUNK_SIZE + fmt.ABYTES + state = sodium.crypto_secretstream_xchacha20poly1305_state() + with encrypted.open('rb') as source, zip_path.open('wb') as target: + source.seek(fmt.HEADER_SIZE) + stream_header = source.read(fmt.STREAM_HEADER_SIZE) + sodium.crypto_secretstream_xchacha20poly1305_init_pull(state, stream_header, key) + last_tag: int | None = None + while chunk := source.read(chunk_size): + plaintext, last_tag = sodium.crypto_secretstream_xchacha20poly1305_pull( + state, chunk + ) + target.write(plaintext) + if last_tag != sodium.crypto_secretstream_xchacha20poly1305_TAG_FINAL: + # A stream that stops without its FINAL tag was cut short. The digest already catches + # truncation of a stored file; this catches a stream that was framed to look complete. + raise BundleGateError(_DECRYPT_FAILURE_MESSAGE, gate=3) + + # ------------------------------------------------------------------ gate 4 + + def _gate_4_inspect(self, zip_path: Path) -> None: + """Inspect the central directory only. Nothing is written to disk by this gate.""" + try: + with zipfile.ZipFile(zip_path) as archive: + entries = archive.infolist() + except zipfile.BadZipFile as error: + raise BundleGateError('bundle content is not a readable archive', gate=4) from error + + if len(entries) > self.limits.max_entries: + raise BundleGateError( + f'archive declares {len(entries)} entries, above the ' + f'{self.limits.max_entries} limit', + gate=4, + ) + if not entries: + raise BundleGateError('archive is empty', gate=4) + + uncompressed = sum(entry.file_size for entry in entries) + compressed = sum(entry.compress_size for entry in entries) + if uncompressed > self.limits.max_uncompressed_bytes: + raise BundleGateError( + f'archive declares {uncompressed} uncompressed bytes, above the ' + f'{self.limits.max_uncompressed_bytes} limit', + gate=4, + ) + ratio = uncompressed / max(compressed, 1) + if ratio > self.limits.max_compression_ratio: + raise BundleGateError( + f'archive compression ratio {ratio:.1f} is above the ' + f'{self.limits.max_compression_ratio} limit', + gate=4, + ) + + for entry in entries: + self._check_entry_name(entry.filename) + self._check_entry_mode(entry) + + @staticmethod + def _check_entry_name(name: str) -> None: + """Reject absolute names and traversal, from the declared name alone.""" + normalised = name.replace('\\', '/') + if normalised.startswith('/') or (len(normalised) > 1 and normalised[1] == ':'): + raise BundleGateError('archive declares an absolute entry path', gate=4) + if any(part == '..' for part in normalised.split('/')): + raise BundleGateError('archive declares an entry escaping its own tree', gate=4) + + @staticmethod + def _check_entry_mode(entry: zipfile.ZipInfo) -> None: + """Reject symlinks and anything that is neither a regular file nor a directory.""" + mode = entry.external_attr >> 16 + if stat.S_IFMT(mode) == 0: + # No file-type bits stored. This is the normal case for two kinds of entry: a zip + # written by a tool that records no Unix mode at all (`mode == 0`), and + # `ZipFile.writestr`, which stores permissions only — the producer's `metadata.json` + # arrives as `0o600 << 16`. Nothing to check: the member is extracted as a plain file + # with the worker's own permissions either way. + return + if stat.S_ISLNK(mode): + raise BundleGateError('archive declares a symbolic link', gate=4) + if not (stat.S_ISREG(mode) or stat.S_ISDIR(mode)): + raise BundleGateError('archive declares an entry that is not a regular file', gate=4) + + # ------------------------------------------------------------------ gate 5 + + def _gate_5_extract(self, zip_path: Path, root: Path) -> None: + """Extract into a fresh directory, containing every member by resolved path. + + A pre-existing extraction root is a failure rather than something to clear: it means + another import, or a previous attempt, owns that path. + """ + if root.exists(): + raise BundleGateError('extraction directory already exists', gate=5) + root.mkdir(parents=True) + resolved_root = root.resolve() + try: + with zipfile.ZipFile(zip_path) as archive: + for entry in archive.infolist(): + self._extract_member(archive, entry, root, resolved_root) + except BundleGateError: + shutil.rmtree(root, ignore_errors=True) + raise + except (OSError, zipfile.BadZipFile) as error: + shutil.rmtree(root, ignore_errors=True) + raise BundleGateError('bundle content could not be unpacked', gate=5) from error + + @staticmethod + def _extract_member( + archive: zipfile.ZipFile, entry: zipfile.ZipInfo, root: Path, resolved_root: Path + ) -> None: + """Write one member, re-checking containment on the resolved path.""" + name = entry.filename.replace('\\', '/') + target = root / name + resolved = Path(os.path.realpath(target)) + if resolved != resolved_root and resolved_root not in resolved.parents: + raise BundleGateError('archive member resolves outside the extraction root', gate=5) + if entry.is_dir(): + resolved.mkdir(parents=True, exist_ok=True) + return + if resolved.exists(): + raise BundleGateError('archive member would overwrite an existing file', gate=5) + resolved.parent.mkdir(parents=True, exist_ok=True) + # Written as a plain file: the archive's stored permissions are never restored, so no + # member can arrive executable. + with archive.open(entry) as source, resolved.open('wb') as sink: + shutil.copyfileobj(source, sink) + + # ------------------------------------------------------------------ gate 6 + + def _gate_6_structure(self, root: Path) -> tuple[dict[str, Any], dict[str, Any]]: + """Validate the extracted layout and the metadata schema.""" + metadata_path = root / METADATA_FILE_NAME + artifacts_dir = root / ARTIFACTS_DIR_NAME + if not metadata_path.is_file(): + raise BundleGateError('bundle has no metadata document at its root', gate=6) + if not artifacts_dir.is_dir(): + raise BundleGateError('bundle has no artifacts directory', gate=6) + + document = self._load_metadata_document(metadata_path) + metadata = document['metadata'] + parameters = document['parameters'] + + self._check_metadata_fields(metadata) + self._check_parameters(parameters) + self._check_model_directories(artifacts_dir) + return metadata, parameters + + @staticmethod + def _load_metadata_document(metadata_path: Path) -> dict[str, Any]: + """Read `metadata.json` and check its two top-level keys.""" + try: + document = json.loads(metadata_path.read_text(encoding='utf-8')) + except (json.JSONDecodeError, UnicodeDecodeError, OSError) as error: + raise BundleGateError('bundle metadata is not readable JSON', gate=6) from error + if not isinstance(document, dict) or set(document) != REQUIRED_TOP_LEVEL_KEYS: + raise BundleGateError( + 'bundle metadata does not carry exactly the parameters and metadata blocks', gate=6 + ) + if not isinstance(document['metadata'], dict) or not isinstance( + document['parameters'], dict + ): + raise BundleGateError('bundle metadata blocks are not objects', gate=6) + return document + + @staticmethod + def _check_metadata_fields(metadata: dict[str, Any]) -> None: + """Every origin field the import records must be present and usable.""" + for field_name in REQUIRED_METADATA_FIELDS: + value = metadata.get(field_name) + if not isinstance(value, str) or not value.strip(): + raise BundleGateError(f'bundle metadata is missing the {field_name} field', gate=6) + experiment_name = metadata['experiment_name'] + if not (EXPERIMENT_NAME_MIN_LENGTH <= len(experiment_name) <= EXPERIMENT_NAME_MAX_LENGTH): + raise BundleGateError( + "bundle experiment name does not fit the platform's limits", gate=6 + ) + + @staticmethod + def _check_parameters(parameters: dict[str, Any]) -> None: + """`target_variable` is what `model_config.target` is built from; nothing else supplies it.""" + for name in REQUIRED_PARAMETERS: + value = parameters.get(name) + if value is None or not str(value).strip(): + raise BundleGateError(f'bundle parameters are missing {name}', gate=6) + + @staticmethod + def _check_model_directories(artifacts_dir: Path) -> None: + """`prediction_model` and `data_model` must each be a loadable MLflow model directory.""" + for directory in REQUIRED_MODEL_DIRS: + model_dir = artifacts_dir / directory + if not model_dir.is_dir(): + raise BundleGateError(f'bundle has no {directory} artifacts', gate=6) + for required in REQUIRED_MODEL_FILES: + if not (model_dir / required).is_file(): + raise BundleGateError(f'bundle {directory} artifacts are incomplete', gate=6) + + # ------------------------------------------------------------------ gate 7 + + def _gate_7_content_policy(self, root: Path) -> None: + """Only the file kinds an MLflow artifact tree calls for may be present.""" + for path in sorted(root.rglob('*')): + if path.is_dir(): + continue + if not path.is_file(): + raise BundleGateError('bundle contains an entry that is not a file', gate=7) + if path.name in ALLOWED_FILE_NAMES: + continue + if path.suffix.lower() in ALLOWED_FILE_SUFFIXES: + continue + raise BundleGateError('bundle contains a file a model export should not carry', gate=7) diff --git a/laborious/utils/bundle/steps.py b/laborious/utils/bundle/steps.py new file mode 100644 index 0000000..35c2506 --- /dev/null +++ b/laborious/utils/bundle/steps.py @@ -0,0 +1,55 @@ +"""The import pipeline's step vocabulary, and the gate-to-step mapping. + +`ImportStep` is the join between three things: the machine-readable step written into the import +record's JSON detail, the key of the error catalog (one code and one sentence per step) and the +value the frontend switches on. It lives in the bundle package rather than in the activities module +because the reader raises with it — `laborious/activities/model_import.py` imports it from here so +there is exactly one definition. + +It holds **pipeline** steps only. The two failures that are not steps of the pipeline — failing to +claim the record and failing to write it — are reported by the failure channel under their own +markers (`not_claimable`, `status_write`), deliberately outside this enum so the +catalog-completeness test keeps meaning "every pipeline step has a sentence". +""" + +from __future__ import annotations + +from enum import StrEnum +from typing import Final + + +class ImportStep(StrEnum): + """Where an import stopped. Values are what gets written to the record.""" + + RECEIVED = 'received' + DOWNLOAD = 'download' + DECRYPTION = 'decryption' + ARCHIVE_INSPECTION = 'archive_inspection' + EXTRACTION = 'extraction' + STRUCTURE_VALIDATION = 'structure_validation' + CONTENT_POLICY = 'content_policy' + EXPERIMENT_CREATION = 'experiment_creation' + ARTIFACT_UPLOAD = 'artifact_upload' + REGISTRATION = 'registration' + MODEL_DOCUMENT = 'model_document' + CLEANUP = 'cleanup' + + +# The seven gates map onto the first seven steps, with gates 1 and 2 sharing one: both say "these +# are not the bytes that were sent" — an object over the size ceiling, a key with the wrong shape, a +# digest that disagrees with what the uploader computed, a wrong magic, an unsupported version, a +# chunk stream that does not match the header's digest. The catalog's `download` sentence ("The +# uploaded file could not be read, or it is not the file that was sent. Upload it again.") is true +# for all of them, and none of them may borrow the `decryption` sentence, which implicates the +# password. +GATE_TO_STEP: Final[dict[int, ImportStep]] = { + 1: ImportStep.DOWNLOAD, + 2: ImportStep.DOWNLOAD, + 3: ImportStep.DECRYPTION, + 4: ImportStep.ARCHIVE_INSPECTION, + 5: ImportStep.EXTRACTION, + 6: ImportStep.STRUCTURE_VALIDATION, + 7: ImportStep.CONTENT_POLICY, +} + +GATE_COUNT: Final[int] = 7 diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py new file mode 100644 index 0000000..aa2a019 --- /dev/null +++ b/laborious/utils/connectors_config.py @@ -0,0 +1,162 @@ +import json +import tempfile +from os import getenv, path +from typing import Any + + +def build_mlflow_config() -> dict[str, Any]: + """ + Build MLFlow server configuration from environment variables. + + This function constructs an MLFlow configuration dictionary from + environment variables with sensible defaults for local development. + It handles server connection and authentication parameters. + + Environment Variables: + MLFLOW_HOST: MLFlow server hostname (default: http://localhost) + MLFLOW_PORT: MLFlow server port (default: 5080) + MLFLOW_USERNAME: MLFlow username (default: aignosi) + MLFLOW_PASSWORD: MLFlow password (default: aignosi) + + Returns: + dict: MLFlow configuration dictionary with all required parameters + """ + return { + 'host': getenv('MLFLOW_HOST', 'http://localhost'), + 'port': int(getenv('MLFLOW_PORT', '5080')), + 'username': getenv('MLFLOW_USERNAME', 'aignosi'), + 'password': getenv('MLFLOW_PASSWORD', 'aignosi'), + } + + +def build_opc_config() -> dict[str, Any]: + """ + Build OPC server configuration from environment variables. + + This function constructs an OPC server configuration dictionary from + environment variables. It supports both single server and multi-server + configurations with flexible parameter handling. + + Environment Variables: + OPC_CONFIG: JSON string containing multiple OPC server configurations + OPC_ID: OPC server ID (fallback, default: 1) + OPC_URL: Single OPC server URL (fallback, default: opc.tcp://localhost:4840) + OPC_SERVER_URI: Single OPC server URI (fallback, default: opc.tcp://localhost:4840) + OPC_CERT_PATH: Client certificate path (fallback, default: None) + OPC_PRIVATE_KEY_PATH: Client private key path (fallback, default: None) + OPC_SERVER_CERT_PATH: Server certificate path (fallback, default: None) + OPC_RECONNECTION_INTERVAL: Reconnection interval in seconds (fallback, default: 120) + + Returns: + dict: OPC server configuration dictionary + """ + opc_raw = getenv('OPC_CONFIG', None) + + if opc_raw: + return json.loads(opc_raw) + + return { + getenv('OPC_ID', '1'): { + 'id': getenv('OPC_ID', '1'), + 'server_name': getenv('OPC_SERVER_NAME', 'default_server'), + 'url': getenv('OPC_URL', 'opc.tcp://localhost:4840'), + 'server_uri': getenv('OPC_SERVER_URI', 'opc.tcp://localhost:4840'), + 'cert_path': getenv('OPC_CERT_PATH', None), + 'private_key_path': getenv('OPC_PRIVATE_KEY_PATH', None), + 'server_cert_path': getenv('OPC_SERVER_CERT_PATH', None), + 'reconnection_interval': int(getenv('OPC_RECONNECTION_INTERVAL', '120')), + } + } + + +def build_minio_config() -> dict[str, Any]: + """ + Build MinIO (S3-compatible) configuration from environment variables. + + Environment Variables: + MINIO_ENDPOINT: MinIO endpoint including scheme (default: http://localhost:9000) + MINIO_ACCESS_KEY: Access key (default: minioadmin) + MINIO_SECRET_KEY: Secret key (default: minioadmin) + MINIO_REGION: Region name for S3 client (default: us-east-1) + MINIO_BUCKET_DEFAULT: Default bucket for uploads (default: laborious) + MINIO_SECURE: Whether to use HTTPS (default: false) + Returns: + dict: MinIO configuration dictionary + """ + return { + 'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'), + 'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'), + 'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'), + 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'), + 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), + 'secure': getenv('MINIO_SECURE', 'false') == 'true', + } + + +def build_import_config() -> dict[str, Any]: + """ + Build the `.sientia` model-import configuration from environment variables. + + The password key is read as the raw base64 string and **not** validated here: the worker must + start whether or not `IMPORT_MODEL_ENABLED` is set, and a missing or malformed key is an import + failure at the decryption step (`laborious.utils.import_password.load_envelope_key`), not a + boot failure. + + Environment Variables: + IMPORT_PASSWORD_KEY: base64 of the 32-byte AES-256-GCM key that wraps bundle passwords + IMPORT_BUNDLE_BUCKET: bucket holding the uploaded bundles (default: sientia) + IMPORT_BUNDLE_PREFIX: key prefix the uploads must sit under (default: imported_models/) + IMPORT_WORK_DIR: worker scratch directory (default: /sientia-import) + IMPORT_MAX_OBJECT_BYTES: gate 1 object size ceiling (default: 1 GiB) + IMPORT_MAX_ARCHIVE_ENTRIES: gate 4 entry count ceiling (default: 5000) + IMPORT_MAX_UNCOMPRESSED_BYTES: gate 4 uncompressed size ceiling (default: 4 GiB) + IMPORT_MAX_COMPRESSION_RATIO: gate 4 compression ratio ceiling (default: 200) + IMPORT_MODELS_COLLECTION: MongoDB collection holding the model listing (default: models) + IMPORT_BUNDLE_RETENTION_DAYS: lifecycle expiry for the uploaded bundle (default: 7) + + Returns: + dict: import configuration dictionary + """ + prefix = getenv('IMPORT_BUNDLE_PREFIX', 'imported_models/') + return { + 'password_key': getenv('IMPORT_PASSWORD_KEY'), + 'bucket': getenv('IMPORT_BUNDLE_BUCKET', 'sientia'), + 'prefix': prefix if prefix.endswith('/') else f'{prefix}/', + 'work_dir': getenv('IMPORT_WORK_DIR', path.join(tempfile.gettempdir(), 'sientia-import')), + 'max_object_bytes': int(getenv('IMPORT_MAX_OBJECT_BYTES', str(1024 * 1024 * 1024))), + 'max_entries': int(getenv('IMPORT_MAX_ARCHIVE_ENTRIES', '5000')), + 'max_uncompressed_bytes': int( + getenv('IMPORT_MAX_UNCOMPRESSED_BYTES', str(4 * 1024 * 1024 * 1024)) + ), + 'max_compression_ratio': float(getenv('IMPORT_MAX_COMPRESSION_RATIO', '200')), + 'models_collection': getenv('IMPORT_MODELS_COLLECTION', 'models'), + 'retention_days': int(getenv('IMPORT_BUNDLE_RETENTION_DAYS', '7')), + } + + +def build_import_status_config() -> dict[str, Any]: + """ + Build the connection to the import log database — a **second** Postgres connection. + + The import log is `public.experiment_run` in the database owned by the Spring Boot BFF + (`sientia-core-mlops-bff`), not the `sientia` database Laborious itself uses. The table and + schema are constants of the code, never inputs; only the connection is configuration. + + Environment Variables: + IMPORT_STATUS_DB_HOST: host (default: localhost) + IMPORT_STATUS_DB_PORT: port (default: 5432) + IMPORT_STATUS_DB_NAME: database (default: sientia-core-mlops-bff) + IMPORT_STATUS_DB_USER: user (default: mlops_bff_user) + IMPORT_STATUS_DB_PASSWORD: password (default: empty) + + Returns: + dict: import status database configuration dictionary + """ + return { + 'host': getenv('IMPORT_STATUS_DB_HOST', 'localhost'), + 'port': int(getenv('IMPORT_STATUS_DB_PORT', '5432')), + 'dbname': getenv('IMPORT_STATUS_DB_NAME', 'sientia-core-mlops-bff'), + 'user': getenv('IMPORT_STATUS_DB_USER', 'mlops_bff_user'), + 'password': getenv('IMPORT_STATUS_DB_PASSWORD', ''), + 'schema': 'public', + } diff --git a/laborious/utils/dataframe_debug.py b/laborious/utils/dataframe_debug.py new file mode 100644 index 0000000..995b606 --- /dev/null +++ b/laborious/utils/dataframe_debug.py @@ -0,0 +1,34 @@ +from typing import Any + +from pandas import DataFrame + +DEFAULT_MAX_DEBUG_DATAFRAME_ROWS = 100 + + +def build_dataframe_debug_message( + message: str, + data: Any, + max_rows: int = DEFAULT_MAX_DEBUG_DATAFRAME_ROWS, +) -> str: + """ + Build a safe debug message for dataframe payloads + + Args: + - message (str): Base message to identify the logged payload + - data (Any): Payload to evaluate for dataframe-aware logging + - max_rows (int): Maximum dataframe row count allowed for full payload logging + + Return: + Formatted debug message with full dataframe content or compact summary + """ + if not isinstance(data, DataFrame): + return f'{message} {data}' + + rows = data.shape[0] + if rows <= max_rows: + return f'{message}\n{data.to_csv()}' + + return ( + f'{message} skipped because dataframe has {rows} rows ' + f'(max: {max_rows}). Shape: {data.shape}' + ) diff --git a/laborious/utils/filters/__init__.py b/laborious/utils/filters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/filters/conditional_filters.py b/laborious/utils/filters/conditional_filters.py new file mode 100644 index 0000000..967e52c --- /dev/null +++ b/laborious/utils/filters/conditional_filters.py @@ -0,0 +1,48 @@ +from pandas import DataFrame + + +def filter_specific_variables_null_values(data: DataFrame, config: dict) -> bool: + """ + Filter to check if specific variables contain null values. + + This function examines a DataFrame to determine if any of the specified variables + contain null (NaN) values. It returns True if null values are found for any of + the specified variables, False otherwise. + + Args: + data (DataFrame): The pandas DataFrame to be examined. Must contain columns + named 'variable' and 'value'. + config (dict): Configuration dictionary containing the following key: + - variables (list): List of variable names to check for null values + + Returns: + bool: True if any of the specified variables contain null values, + False if none of the specified variables contain null values. + + """ + + if data.empty: + return False + + return not data[data['variable'].isin(config['variables']) & data['value'].isna()].empty + + +def filter_empty_data(data: DataFrame, _config: dict) -> bool: + """ + Filter to check if the DataFrame is empty. + + This function determines whether the provided DataFrame contains any data. + It's a simple utility function that can be used in conditional logic to + handle cases where no data is available. + + Args: + data (DataFrame): The pandas DataFrame to be checked for emptiness. + _config (dict): Configuration dictionary (unused in this function). + The underscore prefix indicates this parameter is required for + interface consistency but not used in the implementation. + + Returns: + bool: True if the DataFrame is empty (has no rows), False if it contains data. + + """ + return data.empty diff --git a/laborious/utils/filters/mlflow_filters.py b/laborious/utils/filters/mlflow_filters.py new file mode 100644 index 0000000..82970e8 --- /dev/null +++ b/laborious/utils/filters/mlflow_filters.py @@ -0,0 +1,64 @@ +import numpy as np +from pandas import DataFrame + + +def api_error_filter(response: dict, _config: dict) -> bool: + """ + Filter MLFlow API responses for error conditions. + + This function analyzes MLFlow API responses to detect error conditions + and determine if the response should be filtered out due to quality + or reliability issues. + + + Args: + response: MLFlow API response data (dict) + _config: Filter configuration dictionary + Required keys: + - error_codes (list, optional): List of error codes to detect + - error_keywords (list, optional): List of error keywords to detect + - check_structure (bool, optional): Whether to validate response structure + + Returns: + bool: True if data should be filtered (contains errors), False otherwise + + """ + if not response: + return True + + if not response['success']: + return True + + return False + + +def nan_values_filter(predictions: DataFrame, _config: dict) -> bool: + """ + Filter data for NaN (Not a Number) values. + + This function detects NaN values in MLFlow prediction results and + determines if the data quality is sufficient for further processing + or export operations. + + Args: + predictions: DataFrame containing prediction data to check for NaN values + _config: Filter configuration dictionary + Required keys: + - max_nan_ratio (float, optional): Maximum allowed NaN value ratio (0.0 to 1.0) + - max_nan_count (int, optional): Maximum allowed NaN value count + - check_nested (bool, optional): Whether to check nested data structures + + Returns: + bool: True if data should be filtered (too many NaN values), False otherwise + + """ + data = ( + predictions.replace({None: np.nan}) + .drop(columns=['timestamp'], errors='ignore') + .infer_objects() + ) + + if data.isna().all().all(): + return True + + return False diff --git a/laborious/utils/import_password.py b/laborious/utils/import_password.py new file mode 100644 index 0000000..5b61e8b --- /dev/null +++ b/laborious/utils/import_password.py @@ -0,0 +1,98 @@ +"""The bundle password envelope: AES-256-GCM, opened inside the activity that uses the password. + +The workflow only ever carries `password_envelope`, so no plaintext password reaches the Temporal +event history. This module is imported by the bundle-opening activity and by nothing else — in +particular not by `laborious/workflows/import_model.py`, which a test asserts. + +Envelope layout, chosen because WebCrypto gives the Angular frontend AES-GCM natively: + + base64( nonce(12 bytes) || ciphertext || GCM tag(16 bytes) ) + +The key is a 32-byte value held in a Kubernetes Secret and exposed as `IMPORT_PASSWORD_KEY`, +base64-encoded. Neither the key, the envelope nor the recovered password is ever logged, notified, +used as a metric label or written to the import record. +""" + +from __future__ import annotations + +import base64 +import binascii +from typing import Final + +from cryptography.exceptions import InvalidTag +from cryptography.hazmat.primitives.ciphers.aead import AESGCM + +NONCE_SIZE: Final[int] = 12 +TAG_SIZE: Final[int] = 16 +KEY_SIZE: Final[int] = 32 +MIN_ENVELOPE_SIZE: Final[int] = NONCE_SIZE + TAG_SIZE + + +class EnvelopeError(Exception): + """The envelope or its key is unusable. + + Deliberately one type for every cause — missing key, wrong key, malformed base64, altered + bytes — because the failure is reported at the decryption step either way. What it must never + be confused with is a wrong *bundle* password: that one comes from the reader's gate 3, and the + two carry different sentences. + """ + + +def load_envelope_key(raw_key: str | None) -> bytes: + """Decode the configured envelope key. + + Args: + raw_key: the value of `IMPORT_PASSWORD_KEY` — base64 of exactly 32 bytes. + + Returns: + bytes: the 32-byte key. + + Raises: + EnvelopeError: when the value is absent, not base64, or not 32 bytes long. There is no + fallback interpretation: an unset key never means "the envelope is plaintext". + """ + if not raw_key: + raise EnvelopeError('the import password key is not configured') + try: + key = base64.b64decode(raw_key, validate=True) + except (binascii.Error, ValueError) as error: + raise EnvelopeError('the import password key is not valid base64') from error + if len(key) != KEY_SIZE: + raise EnvelopeError(f'the import password key must be {KEY_SIZE} bytes') + return key + + +def decrypt_password_envelope(envelope: str, key: bytes) -> str: + """Recover the bundle password from its envelope. + + Args: + envelope: base64 of `nonce || ciphertext || tag`. + key: the 32-byte key from `load_envelope_key`. + + Returns: + str: the plaintext password, exactly as it was typed — not trimmed, not case-folded. NFC + normalisation happens at key derivation, where the producer does it. + + Raises: + EnvelopeError: malformed envelope, wrong key or altered bytes (AES-GCM authentication). + """ + if not envelope: + raise EnvelopeError('the import request carries no password envelope') + try: + payload = base64.b64decode(envelope, validate=True) + except (binascii.Error, ValueError) as error: + raise EnvelopeError('the password envelope is not valid base64') from error + if len(payload) < MIN_ENVELOPE_SIZE: + raise EnvelopeError('the password envelope is too short to hold a nonce and a tag') + + nonce, sealed = payload[:NONCE_SIZE], payload[NONCE_SIZE:] + try: + plaintext = AESGCM(key).decrypt(nonce, sealed, None) + except InvalidTag as error: + raise EnvelopeError('the password envelope could not be authenticated') from error + except ValueError as error: + raise EnvelopeError('the password envelope could not be opened') from error + try: + return plaintext.decode('utf-8') + except UnicodeDecodeError as error: + raise EnvelopeError('the recovered password is not valid UTF-8') from error diff --git a/laborious/utils/models/__init__.py b/laborious/utils/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py new file mode 100644 index 0000000..0521f5e --- /dev/null +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -0,0 +1,348 @@ +""" +MinIO-backed DataFrame payload for Temporal workflows. + +Data is never stored as a pandas ``DataFrame`` field on the dataclass. +Instead, the DataFrame is only provided as an input to: +`from_dataframe` / `from_dataframe_to_dict`. + +At build time, the DataFrame is evaluated for its serialized size; if it exceeds +the configured threshold, it is serialized to parquet bytes and uploaded to MinIO. +Otherwise, it is inlined as a Temporal-friendly ``dict``. +""" + +import pickle +import re +from collections.abc import Hashable +from dataclasses import dataclass +from datetime import datetime +from io import BytesIO +from os import getenv +from typing import Any, Literal + +from pandas import DataFrame, read_parquet +from sientia_do.observability.logger import Logger +from sientia_do.repository.minio_repository import MinioRepository +from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now + +# Keys that are part of the serialized wire format (not arbitrary metadata). +_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'}) + +_OBJECT_TIMESTAMP_PATTERN = re.compile( + r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$' +) + +OFFLOAD_THRESHOLD_BYTES = int( + float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024 +) + +# Relative prefix used for storing offloaded prediction datasets in MinIO. +# It is also the root directory for retention cleanup listing. +PREDICTION_DATASETS_PREFIX = 'prediction_datasets' + +OperationKind = Literal['initial', 'transform', 'predict'] + + +def _build_object_key( + model_name: str, operation: OperationKind, timestamp: str +) -> tuple[str, str | None]: + """ + Build the MinIO object key and the directory prefix used for retention listing. + + Args: + model_name: Registered model name used in the pipeline. + operation: Either initial (pre-transform load) or transform (post-MLFlow transform). + timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME. + + Return: + tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root). + """ + # Naming convention: + # - Directory is always `prediction_datasets/` + # - Filename follows the retention-parsing pattern + basename = f'{model_name}-{operation}-{timestamp}.parquet' + model_dir = model_name.strip().strip('/') + prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}' + return f'{prefix}/{basename}', prefix + + +@dataclass +class MinioDataFramePayload: + """ + Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys. + + Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`. + Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class. + """ + + last_timestamp: str + status: dict[str, Any] | None = None + data: dict[Hashable, Any] | None = None + bucket: str | None = None + object_key: str | None = None + object_prefix: str | None = None + uri: str | None = None + + @staticmethod + def _debug( + logger: Logger | None, + message: str, + metadata: dict[str, Any] | None = None, + ) -> None: + """ + Emit debug logs only when logger is provided + + Args: + - logger (Logger | None): Logger instance used for debug messages + - message (str): Message to be logged + - metadata (dict[str, Any] | None): Optional workflow metadata context + """ + if logger is None: + return + logger.custom_debug(message, metadata) + + @classmethod + def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload': + """ + Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization. + + Temporal converts dataclass return values into plain dicts when crossing + workflow/activity boundaries. This method rebuilds the typed instance so + that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are + available on the receiving side. + + If the argument is already a MinioDataFramePayload, it is returned as-is. + + Args: + raw: Dict with keys matching the dataclass fields + (last_timestamp, status, data, bucket, object_key, object_prefix, uri), + or an existing MinioDataFramePayload instance. + + Return: + MinioDataFramePayload: Reconstructed (or original) instance. + """ + if isinstance(raw, MinioDataFramePayload): + return raw + return cls( + last_timestamp=raw['last_timestamp'], + status=raw.get('status'), + data=raw.get('data'), + bucket=raw.get('bucket'), + object_key=raw.get('object_key'), + object_prefix=raw.get('object_prefix'), + uri=raw.get('uri'), + ) + + @staticmethod + def estimate_size_bytes( + df: DataFrame, + metadata: dict[str, Any] | None = None, + logger: Logger | None = None, + ) -> int: + """ + Approximate serialized size of the DataFrame as the default-orient dict. + + Args: + df: DataFrame whose tabular content size is estimated. + + Return: + int: Estimated size in bytes (pickle of dict representation). + """ + try: + size = len(pickle.dumps(df.to_dict())) + except Exception: + size = len(pickle.dumps(df)) + + MinioDataFramePayload._debug( + logger, + f'DataFrame size: {size} bytes', + metadata, + ) + return size + + @staticmethod + def parse_object_timestamp(object_key: str) -> datetime | None: + """ + Parse the timestamp embedded in the object key basename (before .parquet). + + Args: + object_key: S3/MinIO object key whose basename follows + ``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``. + + Return: + datetime | None: Parsed UTC-naive datetime from the key, or None if not matched. + """ + basename = object_key.rsplit('/', 1)[-1] + match = _OBJECT_TIMESTAMP_PATTERN.search(basename) + if not match: + return None + try: + return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME) + except ValueError: + return None + + def cleanup_prefix(self) -> str | None: + """ + Return True if cleanup is enabled for this payload. + """ + if self.object_key is not None and self.data is None: + return self.object_prefix + return None + + def has_data(self) -> bool: + """ + Return True if the payload has some data internally or in MinIO. + """ + return (self.data is not None and self.data != {}) or self.object_key is not None + + @classmethod + async def from_dataframe( + cls, + dataframe: DataFrame | None, + minio_repo: MinioRepository, + model_name: str, + operation: OperationKind, + status: dict[str, Any] | None = None, + workflow_metadata: dict | None = None, + last_timestamp: str | None = None, + logger: Logger | None = None, + ) -> 'MinioDataFramePayload': + """ + Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO. + + The DataFrame is not stored on the returned instance. + + Args: + dataframe: Tabular data to evaluate and persist (inline or MinIO). + metadata: Small metadata dict merged into the payload (e.g. success, message). + minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`. + workflow_metadata: Metadata passed to MinIO store for logging/metrics. + model_name: Registered model name used in the object basename. + operation: Either ``initial`` (query load) or ``transform`` (post-transform). + key_prefix: Backward-compatible parameter (currently ignored for object naming). + size_threshold_bytes: Byte limit before offload. When None, the module-level + environment-derived default is used. + + Return: + MinioDataFramePayload: Instance with data and/or MinIO fields set. + """ + + if dataframe is None or dataframe.empty: + cls._debug( + logger, + 'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload', + workflow_metadata, + ) + return cls( + data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status + ) + + if last_timestamp is None: + last_timestamp = max(dataframe['timestamp'].values.tolist()) + + dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger) + cls._debug( + logger, + ( + f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes ' + f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)' + ), + workflow_metadata, + ) + + if dataframe_size <= OFFLOAD_THRESHOLD_BYTES: + cls._debug( + logger, + 'MinioDataFramePayload.from_dataframe using inline payload', + workflow_metadata, + ) + return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status) + + timestamp = now().strftime(DATETIME_FORMAT_FILENAME) + object_key, object_prefix = _build_object_key(model_name, operation, timestamp) + cls._debug( + logger, + ( + 'MinioDataFramePayload.from_dataframe offloading payload to MinIO ' + f'with key {object_key}' + ), + workflow_metadata, + ) + + # Upload using the relative object key. The upstream repository will + # prefix it internally under its MinIO namespace. + parquet_buffer = BytesIO() + dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True) + file_bytes = parquet_buffer.getvalue() + + upload_result = await minio_repo.upload_file( + file_bytes=file_bytes, + relative_key=object_key, + metadata=workflow_metadata, + ) + + bucket = minio_repo.bucket + object_key_full = upload_result.get('minio_object_name', object_key) + uri = f's3://{bucket}/{object_key_full}' if bucket else None + cls._debug( + logger, + f'MinioDataFramePayload.from_dataframe upload completed: {uri}', + workflow_metadata, + ) + + return cls( + data=None, + bucket=bucket, + object_key=object_key_full, + object_prefix=object_prefix, + uri=uri, + last_timestamp=last_timestamp, + status=status, + ) + + async def retrieve( + self, + minio_repo: MinioRepository, + workflow_metadata: dict[str, Any] | None = None, + logger: Logger | None = None, + ) -> DataFrame: + """ + Load parquet from MinIO when object_key is set and populate inline data. + + Args: + minio_repo: sientia_do MinioRepository (or compatible) with download_file(). + workflow_metadata: Metadata passed to MinIO read for logging/metrics. + + Return: + dict[str, Any]: Flat dict with data filled (same keys as to_dict after load). + """ + if self.data is not None: + self._debug( + logger, + 'MinioDataFramePayload.retrieve using inline payload data', + workflow_metadata, + ) + return DataFrame(self.data) + + if not self.has_data(): + self._debug( + logger, + 'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe', + workflow_metadata, + ) + return DataFrame() + + self._debug( + logger, + f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}', + workflow_metadata, + ) + file_bytes = await minio_repo.download_file( + object_name=self.object_key, metadata=workflow_metadata + ) + df = read_parquet(BytesIO(file_bytes)) + self._debug( + logger, + f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}', + workflow_metadata, + ) + return df diff --git a/laborious/utils/repository/minio_manager.py b/laborious/utils/repository/minio_manager.py new file mode 100644 index 0000000..17c85e6 --- /dev/null +++ b/laborious/utils/repository/minio_manager.py @@ -0,0 +1,32 @@ +from sientia_do.notifications.handlers import NotificationHandler +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 import MinioRepository + + +class MinioManager(SientiaMonitoring): + minio_repository: MinioRepository | None = None + + def __init__( + self, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, + ): + if self.minio_repository is None: + self.minio_repository = minio_repository + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + def close(self) -> None: + """ + Close the MinioManager and clean up resources. + """ + if self.minio_repository is not None: + try: + self.minio_repository.close() + finally: + self.minio_repository = None + + SientiaMonitoring.shutdown(self) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py new file mode 100644 index 0000000..e89b685 --- /dev/null +++ b/laborious/utils/repository/model_repository.py @@ -0,0 +1,1774 @@ +""" +MLflow repository utilities + +This module provides the `MLFlowRepository` class and helpers to interact with +an MLflow tracking server and model registry. It covers model discovery, +downloading/loading with multiple flavors, cached operations with retention +policies, transformation/prediction interfaces, retraining workflows, and +production model promotion. + +Capabilities: +- Model loading and caching with retention policies +- Data transformation and prediction operations +- Model retraining workflows +- Production model updates and versioning +""" + +import ctypes +import gc +import inspect +import threading +import time +import traceback +from collections.abc import Callable +from datetime import datetime, timedelta +from io import StringIO +from os import environ, listdir, makedirs, path +from shutil import rmtree +from typing import Any, Literal, overload + +import joblib +import mlflow +import pandas as pd +from mlflow.entities import Experiment +from numpy import ndarray +from sientia_do.notifications.handlers import NotificationHandler +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 laborious import metrics +from laborious.utils.dataframe_debug import build_dataframe_debug_message + +ARTIFACTS_PATH = './tmp/artifacts' +TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl' +PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl' + +INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'." + + +def force_memory_release(logger: Logger): + """Attempt to release memory from the Python process. + + Executes a garbage collection cycle and calls `malloc_trim(0)` on glibc + where available to return free memory to the OS. This may be a no-op on + non-glibc systems. + + Args: + logger (Logger): Logger for observability. + """ + gc.collect() + + try: + ctypes.CDLL('libc.so.6').malloc_trim(0) + logger.info('Memory released') + except Exception as e: + logger.info(f'Memory release failed: {e}') + + +class MLFlowRepository(SientiaMonitoring): + _MAX_DEBUG_DATAFRAME_ROWS = 100 + + def __init__( + self, + host: str, + username: str, + password: str, + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + ): + """Initialize MLflow client and base state. + + Args: + host (str): MLflow tracking URI. + username (str): MLflow username. + password (str): MLflow password. + logger (Logger): Logger instance. + """ + # set tracking uri + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + mlflow.set_tracking_uri(host) + + environ['MLFLOW_TRACKING_USERNAME'] = username + environ['MLFLOW_TRACKING_PASSWORD'] = password + # Create an MLflow client + self.client = mlflow.tracking.MlflowClient() + + self.model_cache: dict[str, Any] = {} + self._cache_lock = threading.RLock() + self.logger = logger + + 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]): Metadata for contextual logging + """ + self.debug( + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, + ), + metadata, + ) + + """ + Functions related to get model registry parameters + """ + + def get_model_uri(self, run_id: str, prediction: bool = True): + """Build the artifact URI for a run's model. + + Args: + run_id (str): MLflow run identifier. + prediction (bool): If True, return `prediction_model` URI, + otherwise return `data_model` URI. + + Returns: + str: Artifact URI to the selected model within the run. + """ + run_info = mlflow.get_run(run_id) + if prediction: + model_uri = run_info.info.artifact_uri + '/prediction_model' + else: + model_uri = run_info.info.artifact_uri + '/data_model' + return model_uri + + def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str: + """Resolve the run_id for a registered model at a given stage. + + Args: + model_name (str): Registered model name. + stage (str): Desired stage (e.g., 'Production'). + + Returns: + str: Run ID for the latest version at the given stage. + """ + # Use search_registered_models instead of deprecated get_latest_versions + registered_models = self.client.search_registered_models( + filter_string=f"name='{model_name}'" + ) + + if not registered_models: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' not found in the Model Registry." + ) + + # Get all versions of the model and filter by stage + model_versions = self.client.search_model_versions(filter_string=f"name='{model_name}'") + + # Filter versions by the desired stage using current_stage attribute + stage_versions = [mv for mv in model_versions if mv.current_stage == stage] + + if not stage_versions: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' in stage '{stage}' not found in the Model Registry." + ) + + # Sort by version number to get the latest + latest_version = max(stage_versions, key=lambda v: int(v.version)) + source = latest_version.source + if source is None: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' version '{latest_version.version}' in stage '{stage}' " + 'has no source URI to resolve run ID.' + ) + parts = source.split('/') + if len(parts) <= 2 or not parts[2]: + raise mlflow.exceptions.MlflowException( + f"Model '{model_name}' version '{latest_version.version}' in stage '{stage}' " + f"has invalid source URI '{source}' for run ID resolution." + ) + return parts[2] + + def get_next_run_name(self, model_name: str) -> str: + """ + Generate the next run name for a specific MLflow model. + + This method calculates the next sequential run number for a model + by searching existing runs and incrementing the count. It ensures + unique run names for model training and retraining operations. + + Args: + model_name (str): The name of the MLflow model + + Returns: + str: The next run name in format 'model_name-run_number' + """ + runs = mlflow.search_runs(experiment_names=[model_name], order_by=['start_time desc']) + next_run_number = len(runs) + 1 + return f'{model_name}-{next_run_number}' + + def get_experiment( + self, experiment_name: str, create_if_not_exists: bool = False + ) -> Experiment: + """ + Retrieve MLflow experiment by name, optionally creating it. + + This method searches for an MLFlow experiment by name and + returns its unique identifier. It provides error handling + for non-existent experiments. + + Args: + experiment_name (str): Name of the MLflow experiment + + Returns: + Experiment: MLflow experiment object + + Raises: + ValueError: If the experiment name is not found and creation is disabled + """ + experiment = mlflow.get_experiment_by_name(experiment_name) + + if experiment is None: + if create_if_not_exists: + experiment_id = mlflow.create_experiment(experiment_name) + experiment = mlflow.get_experiment(experiment_id) + if experiment is None: + raise ValueError( + f'Experiment {experiment_name} not found after creation, unknown reason' + ) + else: + raise ValueError(f'Experiment {experiment_name} not found') + + return experiment + + def get_model_params(self, run_id: str): + """Fetch parameters associated with a given MLflow run. + + Args: + run_id (str): Run identifier to inspect. + + Returns: + dict: Mapping of parameter names to values. + """ + run_info = mlflow.get_run(run_id) + return run_info.data.params + + def check_artifact_exists( + self, run_id: str, artifact_path: str, metadata: dict[str, Any] + ) -> bool: + """ + Check if an artifact exists in the MLflow Model Registry. + + Args: + run_id (str): Run identifier to inspect. + artifact_path (str): Path to the artifact to check. + + Returns: + bool: True if the artifact exists, False otherwise. + """ + artifacts = self.client.list_artifacts(run_id) + + self.debug(f'Artifacts of {run_id}: \n{artifacts}', metadata) + self.debug(f'Looking for artifact {artifact_path} in {run_id}', metadata) + + return any(artifact.path == artifact_path for artifact in artifacts) + + """ + Functions related to download and load models + """ + + async def dowload_artifacts( + self, model_name: str, metadata: dict[str, Any], artifact_path: str = 'data_model' + ) -> str: + """ + Download artifacts from the latest production run of a model. + + Args: + model_name (str): Registered model name. + metadata (dict[str, Any]): Metadata used for structured logging. + artifact_path (str): Relative path to artifacts within the run. + + Returns: + str: Local filesystem path where artifacts are saved. + """ + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + output_dir = f'{ARTIFACTS_PATH}/{model_name}' + + full_path = path.join(output_dir, artifact_path) + + if path.exists(full_path): + # Remove the directory and create a new one + rmtree(full_path) + makedirs(output_dir, exist_ok=True) + + self.info(f'Downloading artifacts from {run_id} to {output_dir}') + + core_labels = self.get_core_labels(metadata, operation_type='download_artifacts') + + start_time = time.time() + try: + artifacts = self.client.download_artifacts(run_id, artifact_path, output_dir) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + + return artifacts + + async def load_artifact_dataframe( + self, model_name: str, artifact_path: str, metadata: dict[str, Any] + ) -> pd.DataFrame | None: + """ + Load the dataframe content of an artifact from the MLflow Model Registry. + + Args: + model_name (str): The name of the model to download from the registry. + artifact_path (str): The path to the artifact to load. + metadata (dict[str, Any]): Metadata used for structured logging. + + Returns: + pd.DataFrame: The dataframe content of the artifact. + """ + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + core_labels = self.get_core_labels(metadata, operation_type='load_text') + + if not self.check_artifact_exists(run_id, artifact_path, metadata): + return None + + artifact_path = path.join('runs:/', run_id, artifact_path) + + start_time = time.time() + try: + content = mlflow.artifacts.load_text(artifact_path) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + + self.debug(f'Content of {run_id}/{artifact_path}: \n{content}', metadata) + + dataframe = pd.read_csv(StringIO(content)) + + self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata) + return dataframe + + def _find_model_pickle( + self, local_dir: str, artifact_name: str, flavor_error: Exception + ) -> str: + """ + Locate the pickle file of a downloaded model artifact. + + Fallback path only: this is reached after the MLflow flavor loader has + already failed. The target artifact is always a single `.pkl` at the top + level of the artifact directory, so the listing does not recurse — + nested wrapper artifacts must never be deserialized in place of the + model. No other extension is considered and `MLmodel` is not parsed. + + Args: + local_dir (str): Local directory holding the downloaded artifact. + artifact_name (str): Artifact name, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader, + reported in the message so the original cause is not lost. + + Returns: + str: Path to the first top-level `.pkl` file, in name order. + + Raises: + ValueError: If the top level of `local_dir` has no `.pkl` file. + """ + entries = sorted(listdir(local_dir)) + pkl_files = [ + entry + for entry in entries + if entry.endswith('.pkl') and path.isfile(path.join(local_dir, entry)) + ] + + if not pkl_files: + raise ValueError( + f'No .pkl file found in {artifact_name} artifact. ' + f'Original error: {flavor_error}. ' + f'Top-level entries in {local_dir}: {entries}' + ) + + return path.join(local_dir, pkl_files[0]) + + def _load_model_pickle( + self, + local_dir: str, + artifact_name: str, + flavor_error: Exception, + metadata: dict[str, Any] | None = None, + ) -> Any: + """ + Deserialize a downloaded model artifact with `joblib`. + + Fallback path only. `joblib.load` reads the compressed `.pkl` produced + by `joblib.dump`, which is what the MLflow flavor loader fails to read. + The deserialized object is returned as is: no `._model_impl.python_model` + unwrap is applied, because the pickled object already is the model. + + Args: + local_dir (str): Local directory holding the downloaded artifact. + artifact_name (str): Artifact name, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader. + metadata (dict[str, Any] | None): Metadata used for structured + logging, so the chosen pickle path correlates with the other + fallback log lines. + + Returns: + Any: The object deserialized from the artifact's pickle file. + + Raises: + ValueError: If the artifact directory has no top-level `.pkl` file. + Exception: Any error raised by `joblib.load` itself. + """ + pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error) + + self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata) + + return joblib.load(pkl_path) + + async def _download_and_load_pickle( + self, + model_name: str, + artifact_name: str, + flavor_error: Exception, + metadata: dict[str, Any], + run_id: str | None = None, + ) -> Any: + """ + Download a model artifact and deserialize it with `joblib`. + + Fallback path only. The download runs without `dst_path`, so MLflow + writes to a per-call temporary directory: the shared `ARTIFACTS_PATH` + tree is never written to and `dowload_artifacts` — which removes its + destination before downloading — is deliberately not reused. The target + artifact is always a `.pkl` and the deserialized object is returned as + is. + + Callers that already resolved a run to build the flavor URI must pass it + as `run_id`: resolving the Production run again here would open a window + where a promotion between the two lookups makes the fallback serve a + different version than the one the flavor attempt tried and logged. Only + callers whose flavor attempt used a symbolic URI — `models://production`, + resolved by MLflow itself — leave it as `None`. + + Args: + model_name (str): Registered model name. + artifact_name (str): Artifact to download, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader. + metadata (dict[str, Any]): Metadata used for structured logging. + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved here. + + Returns: + Any: The object deserialized from the artifact's pickle file. + + Raises: + ValueError: If the downloaded artifact has no top-level `.pkl` file. + Exception: Any error raised while resolving the Production run, + downloading the artifact, or deserializing the pickle. + """ + if run_id is None: + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + + self.info( + f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}', + metadata, + ) + + local_dir = self.client.download_artifacts(run_id, artifact_name) + + return self._load_model_pickle(local_dir, artifact_name, flavor_error, metadata) + + async def _load_with_joblib_fallback( + self, + primary_loader: Callable[[], Any], + fallback_loader: Callable[[Exception], Any], + failure_description: str, + metadata: dict[str, Any], + fallback_labels: dict[str, Any], + core_labels: dict[str, Any] | None = None, + ) -> Any: + """ + Run a model load with the joblib fallback and its observability. + + Single owner of the fallback flow shared by every load path: warn on the + flavor failure without emitting an error metric, try the fallback, and + emit lag/count under the fallback labels — or the error count chained to + the flavor error when the fallback also fails. Fixes to the flow belong + here, not in the call sites. + + Both loaders are called with no observability concern of their own: + `fallback_loader` receives the flavor error and may be sync or async, + which is what lets the download-based paths and the already-downloaded + wrapper path share this method. + + Args: + primary_loader (Callable[[], Any]): Flavor load attempt. Any + exception it raises triggers the fallback. + fallback_loader (Callable[[Exception], Any]): joblib fallback, + receiving the flavor error. Its result is awaited when + awaitable. + failure_description (str): What failed, e.g. 'Flavor loader sklearn + failed for prediction model x'. The fallback and original error + are appended to it. + metadata (dict[str, Any]): Metadata used for structured logging. + fallback_labels (dict[str, Any]): Metric labels for the fallback + path, already carrying its `operation_type`. + core_labels (dict[str, Any] | None): Metric labels for a successful + flavor load. `None` emits no metric on that path. + + Returns: + Any: The model loaded by the flavor loader or by the fallback. + + Raises: + Exception: The fallback error, chained from the flavor error, when + both attempts fail. + """ + start_time = time.time() + try: + model = primary_loader() + except Exception as flavor_error: + self.warning( + f'{failure_description}, trying joblib fallback. Original error: {flavor_error}', + metadata, + ) + try: + model = fallback_loader(flavor_error) + if inspect.isawaitable(model): + model = await model + except Exception as fallback_error: + await self.emit_metric( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels + ) + raise fallback_error from flavor_error + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + return model + + if core_labels is not None: + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + + return model + + async def load_predict_model( + self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' + ) -> Any: + """ + Load a predictive model from the MLflow Model Registry. + + Args: + model_name (str): The name of the model to download from the registry. + metadata (dict[str, Any]): Metadata used for structured logging. + flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch') + artifact_path (str | None): Path to compressed artifacts if model is compressed + + Returns: + mlflow.pyfunc.PyFuncModel: The loaded predictive model. + + Notes: + - The model is fetched from the "Production" stage of the MLflow Model Registry. + - Warnings during the model loading process are suppressed. + """ + if flavor not in ('pyfunc', 'sklearn', 'pytorch'): + raise ValueError(INVALID_FLAVOR_MESSAGE) + + model_uri = f'models:/{model_name}/production' + self.info(f'Loading prediction model {model_name} from {model_uri}') + + core_labels = self.get_core_labels(metadata, operation_type='load_predict_model') + fallback_labels = self.get_core_labels( + metadata, operation_type='load_predict_model_joblib_fallback' + ) + + def load_flavor() -> Any: + if flavor == 'pyfunc': + return mlflow.pyfunc.load_model(model_uri) + if flavor == 'sklearn': + return mlflow.sklearn.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) + + return await self._load_with_joblib_fallback( + primary_loader=load_flavor, + fallback_loader=lambda flavor_error: self._download_and_load_pickle( + model_name, 'prediction_model', flavor_error, metadata + ), + failure_description=( + f'Flavor loader {flavor} failed for prediction model {model_name}' + ), + metadata=metadata, + fallback_labels=fallback_labels, + core_labels=core_labels, + ) + + async def load_transform_model( + self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' + ) -> Any: + """ + Load the latest Production version of a transformation model. + + This method retrieves the latest production model run ID for the given + model name, constructs the model URI, and loads the model using MLflow. + + Args: + model_name (str): The name of the model to download. + metadata (dict[str, Any]): Metadata used for structured logging. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + artifact_path (str | None): Path to compressed artifacts if model is compressed + + Returns: + Any: The loaded model object, depending on the flavor used. + + Raises: + Exception: If the model run ID or URI cannot be retrieved, or if the + model cannot be loaded. + + Notes: + - The Production run is resolved once and reused by the joblib + fallback, so a promotion between both steps cannot make the + fallback serve a different version than the flavor attempt. + """ + + if flavor not in ('sklearn', 'pyfunc', 'pytorch'): + raise ValueError(INVALID_FLAVOR_MESSAGE) + + latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production') + model_uri = self.get_model_uri(latest_production_id, prediction=False) + + self.info(f'Loading data model {model_name} from {model_uri}') + + core_labels = self.get_core_labels(metadata, operation_type='load_transform_model') + fallback_labels = self.get_core_labels( + metadata, operation_type='load_transform_model_joblib_fallback' + ) + + def load_flavor() -> Any: + if flavor == 'sklearn': + return mlflow.sklearn.load_model(model_uri) + if flavor == 'pyfunc': + return mlflow.pyfunc.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) + + return await self._load_with_joblib_fallback( + primary_loader=load_flavor, + fallback_loader=lambda flavor_error: self._download_and_load_pickle( + model_name, + 'data_model', + flavor_error, + metadata, + run_id=latest_production_id, + ), + failure_description=f'Flavor loader {flavor} failed for data model {model_name}', + metadata=metadata, + fallback_labels=fallback_labels, + core_labels=core_labels, + ) + + async def download_model( + self, + model_name: str, + metadata: dict[str, Any], + model_type: str, + flavor: str, + load_wrapper: bool = False, + ) -> tuple[Any, str | None]: + """ + Download model based on type ("predict" or "transform"). + + Args: + model_name (str): Name of the model to download + metadata (dict[str, Any]): Metadata used for structured logging. + model_type (str): Type of model ('predict' or 'transform') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + load_wrapper (bool): Whether to load wrapper + + Returns: + tuple[Any, str | None]: Model object and optional artifact path. + """ + + self.info( + f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}' + ) + + if model_type not in ['predict', 'transform']: + raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") + + artifact_path = None + + if load_wrapper: + self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}') + + target = 'prediction_model' if model_type == 'predict' else 'data_model' + + artifact_path = await self.dowload_artifacts(model_name, metadata, target) + + self.info( + f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' + ) + + fallback_labels = self.get_core_labels( + metadata, operation_type='download_model_wrapper_joblib_fallback' + ) + + def load_wrapper_flavor() -> Any: + raw_model = mlflow.pyfunc.load_model(artifact_path) + return raw_model._model_impl.python_model + + model = await self._load_with_joblib_fallback( + primary_loader=load_wrapper_flavor, + fallback_loader=lambda flavor_error: self._load_model_pickle( + artifact_path, target, flavor_error, metadata + ), + failure_description=f'Wrapper loader failed for {model_type} model {model_name}', + metadata=metadata, + fallback_labels=fallback_labels, + ) + + self.debug( + f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata + ) + else: + if model_type == 'predict': + model = await self.load_predict_model(model_name, metadata, flavor) + + else: + model = await self.load_transform_model(model_name, metadata, flavor) + + return model, artifact_path + + """ + Functions related to data format + """ + + def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: + """ + Normalize DataFrame index to the expected timestamp string format. + + The index must be timestamp-like. If the index is: + - string: it must match `DATETIME_FORMAT_WITH_TZ` + - datetime or pandas Timestamp: it will be converted to that format + Any other type raises a ValueError. + + Args: + data (pd.DataFrame): DataFrame with timestamp index. + metadata (dict): Metadata for structured logging. + + Returns: + pd.DataFrame: DataFrame with converted datetime index. + """ + if data.empty: + self.info('Data is empty, skipping datetime index detection and parsing', metadata) + return data + + index = data.index + + # Get type of first element of 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, string in format {DATETIME_FORMAT_WITH_TZ}.' + + # Check if all in index are of the same type + 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)}') + + # Check type and converts to DATETIME_FORMAT_WITH_TZ + if index_type is str: + # Validate format of string and return error if not valid + 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 == datetime or index_type == pd.Timestamp: + index = data.index + if hasattr(index, 'tz') and index.tz is None: + data.index = index.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 + + """ + Functions related to cache management of models + """ + + def check_cache_retention(self, cache: dict, retention: int) -> bool: + """ + Check whether cached model data is still valid. + + Args: + cache (dict): Cached model data with a 'timestamp' key. + retention (int): Retention time in minutes. + + Returns: + bool: True if cache is still valid, False if expired. + """ + current_time = datetime.now() + cache_time = cache['timestamp'] + if current_time - cache_time >= timedelta(minutes=retention): + return False + return True + + def handle_valid_model(self, model_name: str, cache: dict) -> dict: + """ + Return the cached model configuration when retention is valid. + + Args: + model_name (str): Name of the model (for logging/consistency). + cache (dict): Cached model data structure. + + Returns: + dict: Model configuration. + """ + self.debug(f'Model {model_name} is still valid, using cached version') + + return cache['target'] + + def handle_outdated_model(self, model_name: str, model_key: str) -> None: + """ + Clean up outdated cached model and its artifacts. + + Args: + model_name (str): Name of the model. + model_key (str): Cache key for the model. + + Returns: + None + """ + self.debug(f'Model {model_name} is outdated, downloading a new one') + + del self.model_cache[model_key]['target'] + del self.model_cache[model_key] + + async def get_model( + self, + model_name: str, + metadata: dict[str, Any], + retention: int, + model_type: str, + flavor: str, + ) -> Any: + """ + Retrieve a model with caching support based on retention policy. + + Args: + model_name (str): Name of the model to retrieve + metadata (dict[str, Any]): Metadata used for structured logging. + retention (int): Cache retention time in minutes (0 = no cache). + model_type (str): Type of model ('predict' or 'transform') + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + + Returns: + Any: Model object. + """ + # Retention is 0, download a new model + if retention <= 0: + model, _artifact_path = await self.download_model( + model_name=model_name, + metadata=metadata, + model_type=model_type, + flavor=flavor, + load_wrapper=False, + ) + return model + + model_key = f'{model_name}_{model_type}' + + # Acquire lock to check cache + with self._cache_lock: + if model_key in self.model_cache: + cache = self.model_cache[model_key] + + # Check if config has changed or is outdated + if self.check_cache_retention(cache, retention): + return self.handle_valid_model(model_name=model_name, cache=cache) + else: + # Model is outdated, delete old model files + self.handle_outdated_model(model_name=model_name, model_key=model_key) + else: + self.debug( + f'Model {model_name} is not in {model_type} cache, downloading a new one' + ) + + # Donwload new model (without lock to avoid blocking other threads) + model, _artifact_path = await self.download_model( + model_name=model_name, + metadata=metadata, + model_type=model_type, + flavor=flavor, + load_wrapper=False, + ) + + # Update cache with lock + with self._cache_lock: + cache = {'target': model, 'timestamp': datetime.now()} + self.model_cache[model_key] = cache + + return model + + @overload + async def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['transform'], + retention: int, + flavor: str, + metadata: dict[str, Any], + ) -> pd.DataFrame: ... + + @overload + async def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: Literal['predict'], + retention: int, + flavor: str, + metadata: dict[str, Any], + ) -> pd.DataFrame | ndarray: ... + + async def get_cached_operation( + self, + model_name: str, + data: pd.DataFrame, + operation: str, + retention: int, + flavor: str, + metadata: dict[str, Any], + ) -> pd.DataFrame | ndarray: + """ + Execute a cached operation using the requested model. + + Args: + model_name (str): Registered model name. + data (pd.DataFrame): Input data. + retention (int): Cache retention in minutes. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch'). + metadata (dict[str, Any]): Metadata used for structured logging. + Returns: + pd.DataFrame | ndarray: Operation result. + """ + if operation not in ['transform', 'predict']: + raise ValueError("Invalid operation. Use 'transform' or 'predict'.") + + model = await self.get_model( + model_name=model_name, + metadata=metadata, + retention=retention, + model_type=operation, + flavor=flavor, + ) + + prediction = model.predict(data) + + if retention == 0: + self.info(f'Deleting model {model_name}:{operation} from memory') + del model + + force_memory_release(self.logger) + + return prediction + + """ + Functions related to model retraining + """ + + def get_prediction_data( + self, + prediction_model: Any, + retrain_dataset: pd.DataFrame, + target_name: str, + predict_flavor: str, + ) -> pd.DataFrame: + """ + Get prediction data from prediction model. + """ + input_index = retrain_dataset.index + + if predict_flavor == 'pyfunc': + prediction_data = prediction_model.predict({}, retrain_dataset) + else: + prediction_data = prediction_model.predict(retrain_dataset) + + if isinstance(prediction_data, pd.DataFrame): + prediction_data.columns = pd.Index(['prediction']) + + else: + prediction_data = pd.DataFrame(prediction_data, columns=['prediction']) + + prediction_data.index = input_index + + # Merge prediction data with retrain_dataset on index + prediction_data = pd.merge( # NOSONAR + retrain_dataset, prediction_data, left_index=True, right_index=True, how='left' + ) + + # Rename column "target_name" to "target" + prediction_data.rename(columns={target_name: 'target'}, inplace=True) + + prediction_data['timestamp'] = prediction_data.index + + prediction_data.reset_index(drop=True, inplace=True) + + prediction_data = prediction_data.sort_values(by='timestamp', ascending=True) + + return prediction_data + + async def fit_models( + self, + model_name: str, + data: pd.DataFrame, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + skip_transform: bool = False, + predict_flavor: str = 'sklearn', + target_name: str | None = None, + ) -> dict[str, Any]: + """ + Prepare models and data for a retraining run. + + This method sets up the complete environment for model retraining by: + 1. Loading the current production prediction model + 2. Loading the current production transformation model + 3. Fitting the transformation model with new data + 4. Preparing data for prediction model retraining + 5. Setting up the MLFlow experiment context + + Args: + model_name (str): Name of the MLflow model to retrain. + data (pd.DataFrame): Training data for model retraining. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + target_name (str | None): Optional target column; if None, use model target. + metadata (dict): Metadata for logging. + + Returns: + dict[str, dict[str, Any]]: Mapping with prepared `prediction_model` and + `data_model`, including optional artifact paths. + """ + + self.info(f'Starting model experiment creation for {model_name}', metadata) + self.debug( + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, + ) + + # data.to_csv( + # f"tmp/retrain_data_{model_name}.csv", index=True) + + self.info(f'Retrieved latest production run ID: {latest_production_id}', metadata) + self.info(f'Loading transformation model for {model_name}', metadata) + + load_transform_wrapper = transform_flavor == 'pyfunc' + + data_model, data_artifact_path = await self.download_model( + model_name=model_name, + metadata=metadata, + model_type='transform', + flavor=transform_flavor, + load_wrapper=load_transform_wrapper, + ) + + self.info(f'Loading prediction model for {model_name}', metadata) + + load_predict_wrapper = predict_flavor == 'pyfunc' + + prediction_model, prediction_artifact_path = await self.download_model( + model_name=model_name, + metadata=metadata, + model_type='predict', + flavor=predict_flavor, + load_wrapper=load_predict_wrapper, + ) + + if not skip_transform: + treated_data_candidate = data_model.fit(data) + else: + treated_data_candidate = data_model + + if not isinstance(treated_data_candidate, pd.DataFrame): + data_model = treated_data_candidate + treated_data = data_model.predict(data) + else: + treated_data = treated_data_candidate + + # Stores current index as timestamp, Courier model expects a timestamp column + # with specific format + treated_data['timestamp'] = treated_data.index + + # Parses timestamp column to datetime format to align with data + treated_data = self.detect_and_parse_datetime_index(treated_data, metadata) + + treated_data = treated_data.drop_duplicates(subset=['timestamp'], keep='first') + + self.debug(f'Treated data index: {treated_data.index}', metadata) + + # treated_data.to_csv( + # f"tmp/retrain_treated_data_{model_name}.csv", index=True) + + self.debug(f'Transformed data shape: {treated_data.shape}', metadata) + + if target_name is None: + target_name = data_model.target_variable + self.debug(f'Using target variable from data model: {target_name}', metadata) + else: + self.debug(f'Using provided target variable: {target_name}', metadata) + + # Check if treated_data contains target variable + if target_name not in treated_data.columns: + self.debug( + f'Target variable {target_name} not found in treated data, aligning data with treated data indexes', + metadata, + ) + # Aligns data with treated data indexes to get target variable + aligned_data = data.loc[treated_data.index] + aligned_series = aligned_data[target_name] + retrain_dataset = pd.merge( # NOSONAR + treated_data, aligned_series, left_index=True, right_index=True + ) + else: + # Uses target variable from treated data + self.debug(f'Target variable {target_name} found in treated data, using it', metadata) + retrain_dataset = treated_data + + # retrain_dataset.to_csv( + # f"tmp/retrain_retrain_dataset_{model_name}.csv", index=True) + + prediction_model.fit(retrain_dataset) + + # get prediction data + prediction_data = self.get_prediction_data( + prediction_model, retrain_dataset, target_name, predict_flavor + ) + + self.info(f'Model experiment creation completed successfully for {model_name}', metadata) + + retrain_data = { + 'prediction_model': { + 'model': prediction_model, + 'artifact_path': prediction_artifact_path, + }, + 'data_model': {'model': data_model, 'artifact_path': data_artifact_path}, + 'prediction_data': prediction_data, + } + return retrain_data + + def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]: + """Resolve the code dependencies to forward when re-logging a pyfunc model. + + MLflow flattens whatever was given as `code_paths` into `/code/`, + so the layout differs per model: some ship a `utils` package, others a + single self-contained module (`_model.py`). Re-logging must forward + the entries that are actually on disk — assuming one fixed layout makes + `_validate_and_copy_code_paths` raise `FileNotFoundError` for the other. + + Args: + artifact_path (str | None): Local directory of the downloaded model + artifact, e.g. `./tmp/artifacts//prediction_model`. + metadata (dict): Metadata for structured logging. + + Returns: + list[str]: Existing paths inside `/code`, empty when + the model carries no code directory. + """ + if not artifact_path: + self.warning('No artifact path available, logging model without code paths', metadata) + return [] + + code_dir = path.join(artifact_path, 'code') + + if not path.isdir(code_dir): + self.warning( + f'No code directory found at {code_dir}, logging model without code paths', + metadata, + ) + return [] + + return [ + path.join(code_dir, entry) + for entry in sorted(listdir(code_dir)) + if entry != '__pycache__' + ] + + async def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): + """Log a model into the active MLflow run. + + Args: + model_data (dict): Model holder with keys 'model' and optional 'artifact_path'. + flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch'). + model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'. + metadata (dict): Metadata for structured logging. + """ + model = model_data['model'] + + self.debug(f'Logging {model_type} model to {model_type}', metadata) + + core_labels = self.get_core_labels(metadata, operation_type='log_model') + start_time = time.time() + + try: + if flavor == 'sklearn': + mlflow.sklearn.log_model(model, model_type) + elif flavor == 'pyfunc': + code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata) + + self.debug(f'Code path: {code_path}', metadata) + + model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False) + + self.debug('Model uploaded successfully', metadata) + elif flavor == 'pytorch': + mlflow.pytorch.log_model(model, model_type) + else: + raise ValueError(INVALID_FLAVOR_MESSAGE) + + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels) + + async def create_new_experiment( + self, + model_name: str, + data: pd.DataFrame, + retrain_data: dict, + latest_production_id: str, + metadata: dict, + transform_flavor: str = 'sklearn', + predict_flavor: str = 'sklearn', + ) -> dict: + """ + Execute the complete model retraining process in MLflow. + + This method performs the actual model retraining by: + 1. Starting a new MLFlow run with descriptive metadata + 2. Logging model parameters and hyperparameters + 3. Retraining both prediction and transformation models + 4. Logging training data as artifacts + 5. Saving retrained models to MLFlow registry + + Args: + prediction_model: MLFlow prediction model to retrain + data_model: MLFlow transformation model to retrain + experiment (str): MLflow experiment name for the retraining. + model_name (str): Name of the model being retrained. + data (pd.DataFrame): Training data used for retraining. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + metadata (dict): Metadata for logging. + + Returns: + dict: Metadata about the created run and experiment. + """ + + prediction_model = retrain_data['prediction_model'] + data_model = retrain_data['data_model'] + prediction_data = retrain_data['prediction_data'] + + model_temp_path = path.join(ARTIFACTS_PATH, model_name) + + self.info(f'Starting model retraining process for {model_name}', metadata) + + original_params = self.get_model_params(latest_production_id) + retrain_params = { + **original_params, + 'retrain': True, + 'retrain_date': datetime.now().isoformat(), + 'source_run_id': latest_production_id, + 'retrain_samples': str(data.shape), + } + experiment_description = f'Retrain model {model_name} with new data' + + experiment = self.get_experiment(model_name, create_if_not_exists=True) + experiment_name = experiment.name + + current_run_name = self.get_next_run_name(experiment_name) + + self.debug(f'Attributes: {retrain_params}', metadata) + + data_path = f'{model_temp_path}/retrain_data.csv' + prediction_data_path = f'{model_temp_path}/evaluation_data.csv' + + makedirs(model_temp_path, exist_ok=True) + + data.to_csv(data_path, index=False) + prediction_data.to_csv(prediction_data_path, index=False) + + self.info( + f'Starting model upload for {experiment_name} with run name {current_run_name}', + metadata, + ) + + core_labels = self.get_core_labels(metadata, operation_type='create_new_experiment') + start_time = time.time() + try: + with mlflow.start_run( + experiment_id=experiment.experiment_id, + run_name=current_run_name, + description=experiment_description, + ) as _run: + run_id = _run.info.run_id + self.info('Logging data model', metadata) + # dynamic parameters, including model itself + await self.log_model(data_model, transform_flavor, 'data_model', metadata) + + # dynamic parameters, including model itself + self.info('Logging prediction model', metadata) + await self.log_model(prediction_model, predict_flavor, 'prediction_model', metadata) + + self.info(f'Model logged successfully for {model_name}', metadata) + + self.info(f'Logging remaining parameters for {model_name}', metadata) + + # update transfomation model + # fixed parameters + mlflow.log_params(retrain_params) + + # log the data raw + mlflow.log_artifact(data_path) + mlflow.log_artifact(prediction_data_path) + + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels) + + self.info('Deleting model from filesystem', metadata) + if path.exists(model_temp_path): + rmtree(model_temp_path) + + self.info('Deleting prediction model from memory', metadata) + del prediction_model['model'] + del prediction_model + + self.info('Deleting data model from memory', metadata) + del data_model['model'] + del data_model + + force_memory_release(self.logger) + + return { + 'run_id': run_id, + 'experiment_id': experiment.experiment_id, + 'experiment_name': experiment.name, + } + + async def update_production_model_by_run_id( + self, run_id: str, model_name: str, metadata: dict + ) -> dict: + """ + Promote a specific run's model to Production. + + This method promotes a model from a specific MLFlow run to + production stage. It handles model registration, versioning, + and stage transitions with proper error handling. + + Args: + run_id (str): MLflow run ID containing the model to promote. + model_name (str): Name of the MLflow model. + metadata (dict): Metadata for logging. + + Returns: + dict: Model update metadata containing: + - model_name (str): Name of the updated model + - version (str): New model version number + - mlflow_run_id (str): Source run ID + + Update Process: + 1. Registers the model from the specified run + 2. Retrieves the latest model version + 3. Transitions the model to 'Production' stage + 4. Archives existing production versions + """ + + self.info( + f'Starting production model update for {model_name} with run ID: {run_id}', metadata + ) + + # Registrar o modelo + # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro. + # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso. + + core_labels = self.get_core_labels(metadata, operation_type='register_model') + start_time = time.time() + try: + mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels) + + # Obter a versão mais recente registrada do modelo + model_versions = self.client.get_registered_model(model_name).latest_versions + + if not isinstance(model_versions, list): + raise ValueError('Model versions is not a list') + + max_version = max(model_versions, key=lambda x: int(x.version)).version + + # Mover a versão mais recente do modelo para o estágio de 'Production' + core_labels = self.get_core_labels( + metadata, operation_type='transition_model_version_stage' + ) + start_time = time.time() + try: + self.client.transition_model_version_stage( + name=model_name, + version=max_version, + stage='Production', + archive_existing_versions=True, + ) + except Exception as e: + await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels) + raise e + + await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels) + + return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id} + + """ + Functions that provide the interface to model operations + """ + + async def transform( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ) -> dict[str, Any]: + """ + Transform data using a cached transformation model. + + This method provides a high-level interface for data transformation operations + using MLFlow models. It handles model caching, error management, and data + format conversion automatically. + + Process Flow: + 1. Retrieves or downloads the transformation model using caching mechanism + 2. Applies the transformation model to the input data + 3. Converts the transformed data to dictionary format for API response + 4. Handles any exceptions and returns structured error information + 5. Manages model lifecycle based on retention policy (cleanup artifacts if needed) + + Parameters: + model_name (str): The name of the MLflow model to use for transformation. + data (pd.DataFrame): The input data to be transformed by the model. + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Response dictionary containing: + - success (bool): Operation success status + - content (dict): Transformed data as dictionary, or error information + if operation failed. Error content includes: + - message (str): Error description + - traceback (str): Full exception traceback + + Raises: + Exception: Any exception during model loading or transformation is caught + and returned in the response structure rather than propagated. + """ + + self._debug_dataframe('Data received for model transformation:', data, metadata) + + # data.to_csv( + # f"tmp/data_{model_name}.csv", index=True) + + model_retention = model_config.get('retention_minutes', 0) + flavor = model_config.get('transform_flavor', 'sklearn') + + try: + transformed_data: pd.DataFrame = await self.get_cached_operation( + model_name=model_name, + data=data, + operation='transform', + retention=model_retention, + flavor=flavor, + metadata=metadata, + ) + + self._debug_dataframe( + 'Data received from model transformation:', transformed_data, metadata + ) + + # transformed_data.to_csv( + # f"tmp/transformed_data_{model_name}.csv", index=True) + + transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) + + return {'success': True, 'content': transformed_data} + + except Exception as e: + return { + 'success': False, + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, + } + + async def predict( + self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict + ): + """ + Generate predictions using a cached prediction model. + + This method provides a high-level interface for model prediction operations + using MLFlow models. It handles model caching, performance monitoring, + response formatting, and error management automatically. + + Process Flow: + 1. Preserves input data index for result alignment + 2. Records prediction start time for performance measurement + 3. Retrieves or downloads the prediction model using caching mechanism + 4. Executes model prediction on the input data + 5. Formats predictions into DataFrame with proper column naming + 6. Restores original data index to maintain data alignment + 7. Calculates and adds response time measurement + 8. Converts results to dictionary format for API response + 9. Handles any exceptions and returns structured error information + + Parameters: + model_name (str): The name of the MLflow model to use for prediction. + data (pd.DataFrame): The input data to make predictions on. + model_retention (int): Cache retention time in minutes (0 = no caching). + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Response dictionary containing: + - success (bool): Operation success status + - content (dict): Prediction results as dictionary with: + - prediction: Model predictions array + - response_time: Prediction execution time in seconds + Or error information if operation failed: + - message (str): Error description + - traceback (str): Full exception traceback + + Raises: + Exception: Any exception during model loading or prediction is caught + and returned in the response structure rather than propagated. + """ + + model_retention = model_config.get('retention_minutes', 0) + flavor = model_config.get('predict_flavor', 'sklearn') + + try: + input_index = data.index + start_time = datetime.now() + + self._debug_dataframe('Data received for model prediction:', data, metadata) + + # data.to_csv( + # f"tmp/treated_data_{model_name}.csv", index=True) + + predict_data = await self.get_cached_operation( + model_name=model_name, + data=data, + operation='predict', + retention=model_retention, + flavor=flavor, + metadata=metadata, + ) + + end_time = datetime.now() + + if isinstance(predict_data, pd.DataFrame): + self._debug_dataframe( + 'Data received from model prediction:', predict_data, metadata + ) + + # predict_data.to_csv( + # f"tmp/predicted_data_{model_name}.csv", index=True) + 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.to_csv( + # f"tmp/predicted_data_{model_name}.csv", index=True) + + predict_data.index = input_index + predict_data['response_time'] = (end_time - start_time).total_seconds() + + return {'success': True, 'content': predict_data} + + except Exception as e: + return { + 'success': False, + 'content': {'message': str(e), 'traceback': traceback.format_exc()}, + } + + async def retrain_model( + self, data: pd.DataFrame, model_name: str, model_config: dict, metadata: dict + ) -> dict[str, Any]: + """ + Orchestrate the complete model retraining workflow. + + This method coordinates the entire model retraining process by managing + the MLFlow experiment lifecycle, model loading, retraining execution, + and artifact management. It provides a comprehensive retraining solution + that maintains model versioning and experiment tracking. + + Process Flow: + 1. Creates MLFlow experiment environment: + - Loads current production prediction model + - Loads current production transformation model + - Fits transformation model with new training data + - Prepares transformed data for prediction model retraining + - Sets up MLFlow experiment context + 2. Executes model retraining: + - Starts new MLFlow run with descriptive metadata + - Logs model parameters and hyperparameters + - Retrains both prediction and transformation models + - Logs training data as artifacts + - Saves retrained models to MLFlow registry + - Cleans up temporary files + 3. Returns comprehensive retraining results + + Args: + data (pd.DataFrame): Training data for model retraining. Must contain + all features required by both transformation and + prediction models, including target variable. + model_name (str): Name of the MLflow model to retrain. Must exist + in the MLflow Model Registry in Production stage. + model_config (dict): Model configuration parameters + metadata (dict): Metadata for logging + + Returns: + dict: Retraining operation results and experiment details. + + Raises: + mlflow.exceptions.MlflowException: If model not found in registry + ValueError: If experiment cannot be created or models cannot be loaded + Exception: Any other exception during the retraining process + """ + + self.info(f'Starting model retraining workflow for {model_name}', metadata) + self._debug_dataframe('Data received for model retraining:', data, metadata) + + target_name = model_config.get('target', None) + + transform_flavor = model_config.get('transform_flavor', 'sklearn') + predict_flavor = model_config.get('predict_flavor', 'sklearn') + skip_transform = model_config.get('skip_transform', False) + + self.debug( + f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}', + metadata, + ) + + try: + latest_production_id = self.get_model_run_id(model_name, stage='Production') + self.info('Creating model experiment environment', metadata) + retrain_data = await self.fit_models( + model_name=model_name, + data=data, + transform_flavor=transform_flavor, + skip_transform=skip_transform, + predict_flavor=predict_flavor, + target_name=target_name, + metadata=metadata, + latest_production_id=latest_production_id, + ) + self.info(f'Model experiment created successfully: {retrain_data}', metadata) + + self.info('Saving model retrain', metadata) + experiment = await self.create_new_experiment( + model_name=model_name, + data=data, + retrain_data=retrain_data, + transform_flavor=transform_flavor, + predict_flavor=predict_flavor, + metadata=metadata, + latest_production_id=latest_production_id, + ) + self.info( + f'Model retraining completed successfully for experiment: {experiment}', metadata + ) + + return { + 'success': True, + 'experiment': experiment, + 'message': 'Model retrained successfully.', + } + except Exception as 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(), + } + + async def update_production_model( + self, experiment: dict[str, Any], model_name: str, metadata: dict + ) -> dict: + """ + Update production model using the latest retraining run. + + This method orchestrates the complete production model update process by + identifying the most recent retraining run and promoting it to production + stage. It handles model registration, versioning, and stage transitions + with comprehensive metadata tracking. + + Process Flow: + 1. Retrieves experiment information: + - Converts experiment name to MLFlow experiment ID + - Searches for the most recent retraining run in the experiment + - Filters runs by 'retrain' parameter and orders by completion time + 2. Promotes model to production: + - Registers the model from the specified run to MLFlow Model Registry + - Retrieves the latest model version number + - Transitions the model to 'Production' stage + - Archives existing production versions automatically + 3. Returns comprehensive update metadata + + Args: + experiment (str): MLflow experiment name containing the retraining runs. + Must be a valid experiment that exists in MLFlow. + model_name (str): Name of the MLFlow model to update. Must exist + in the MLFlow Model Registry. + metadata (dict): Metadata for logging + + Returns: + dict: Complete model update metadata containing: + - model_name (str): Name of the updated model + - version (str): New model version number (incremented automatically) + - mlflow_run_id (str): Source run ID of the promoted model + - mlflow_experiment_id (int): Experiment ID for tracking + + Raises: + ValueError: If experiment not found or model versions are invalid + mlflow.exceptions.MlflowException: If model registration or stage + transition fails + Exception: Any other exception during the update process + + Note: + This operation is irreversible. The previous production model will + be automatically archived when the new version is promoted. + """ + run_id = experiment['run_id'] + experiment_id = experiment['experiment_id'] + metadata_result = await self.update_production_model_by_run_id(run_id, model_name, metadata) + + metadata_result['mlflow_experiment_id'] = experiment_id + + return metadata_result diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py new file mode 100644 index 0000000..2ecc27a --- /dev/null +++ b/laborious/utils/repository/opc_repository.py @@ -0,0 +1,866 @@ +import asyncio +import json +import time +import traceback +from datetime import datetime +from pathlib import Path +from typing import Any + +from asyncua import Client +from asyncua.crypto.security_policies import SecurityPolicyBasic256 +from asyncua.ua import DataValue, Variant, VariantType +from asyncua.ua.uaerrors import UaStatusCodeError +from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler +from sientia_do.notifications.models import NotificationLevel +from sientia_do.observability.logger import Logger +from sientia_do.observability.metrics_controller import MetricsController +from sientia_do.observability.sientia_monitoring import SientiaMonitoring + +from laborious import metrics + +# Requested session and secure channel lifetime (ms) before server revision; 10 minutes. +OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS = 10 * 60 * 1000 + + +class OpcClientAlreadyExistsError(RuntimeError): + """Raised when _create_client is called while self.client is already set.""" + + +class OpcSessionAlreadyConnectedError(RuntimeError): + """Raised when _open_session is called while a UA session is already open.""" + + +class OpcClientNotInitializedError(RuntimeError): + """Raised when _open_session is called before _create_client.""" + + +RECONNECTABLE_OPC_BAD_NAMES: frozenset[str] = frozenset( + { + 'BadSessionIdInvalid', + 'BadSessionClosed', + 'BadSessionNotActivated', + 'BadSecureChannelIdInvalid', + 'BadSecureChannelClosed', + 'BadSecureChannelTokenUnknown', + 'BadTcpSecureChannelUnknown', + 'BadServerNotConnected', + 'BadConnectionClosed', + 'BadDisconnect', + 'BadConnectionRejected', + 'BadCommunicationError', + 'BadRequestInterrupted', + 'BadUnknownResponse', + 'BadTimeout', + 'BadRequestTimeout', + 'BadSequenceNumberInvalid', + 'BadSequenceNumberUnknown', + 'BadSecurityModeInsufficient', + 'BadRequestHeaderInvalid', + 'BadInvalidState', + } +) + + +def _opc_authentication_token_str(client: Client | None) -> str: + """ + Serialize the current OPC UA authentication token (session handle) for logging and metrics. + + Return: + str: Token string, or "unknown" if unavailable. + """ + if client is None: + return 'unknown' + try: + proto = client.uaclient.protocol + if proto is None: + return 'unknown' + tok = getattr(proto, 'authentication_token', None) + if tok is None: + return 'unknown' + return str(tok) + except Exception: + return 'unknown' + + +def _opc_status_from_exception(exc: BaseException) -> str: + """ + Resolve OPC UA status name from an exception, including chained UaStatusCodeError causes. + + Args: + exc (BaseException): Raised error from asyncua. + + Return: + str: Status class name or generic Python exception name. + """ + current: BaseException | None = exc + while current is not None: + if isinstance(current, UaStatusCodeError): + return type(current).__name__ + current = current.__cause__ + return type(exc).__name__ + + +def is_reconnectable_opcua_bad(exc: BaseException) -> bool: + """ + Return whether the exception is a Tier-1 OPC UA Bad* that should trigger reconnect. + + Args: + exc (BaseException): Raised error from get_node or write_value. + + Return: + bool: True if reconnect should be scheduled. + """ + return _opc_status_from_exception(exc) in RECONNECTABLE_OPC_BAD_NAMES + + +def _model_labels_from_write_metadata(metadata: dict[str, Any] | None) -> dict[str, str]: + """ + Extract model_id and model_name from write metadata for Prometheus labels. + + Args: + metadata (dict[str, Any] | None): Context passed into write_data; may omit keys. + + Return: + dict[str, str]: Labels model_id and model_name, defaulting to "unknown". + """ + if not metadata: + return {'model_id': 'unknown', 'model_name': 'unknown'} + return { + 'model_id': str(metadata.get('model_id', 'unknown')), + 'model_name': str(metadata.get('model_name', 'unknown')), + } + + +data_type_map = { + 'float': { + 'converter': float, + 'opc_type': VariantType.Float, + }, + 'double': { + 'converter': float, + 'opc_type': VariantType.Double, + }, + 'int': { + 'converter': int, + 'opc_type': VariantType.Int32, + }, + 'bool': { + 'converter': bool, + 'opc_type': VariantType.Boolean, + }, + 'str': { + 'converter': str, + 'opc_type': VariantType.String, + }, +} + + +class OpcRepository(SientiaMonitoring): + def __init__( + self, + opc_id: str, + url: str, + server_name: str, + logger: Logger, + notification_handler: NotificationHandler, + metrics_controller: MetricsController, + reconnection_interval: int = 60, + server_uri: str | None = None, + cert_path: str | None = None, + private_key_path: str | None = None, + server_cert_path: str | None = None, + ): + self.url = url + self.id = opc_id + self.server_name = server_name + self.server_uri = server_uri + self.cert_path = cert_path + self.private_key_path = private_key_path + self.server_cert_path = server_cert_path + self.reconnection_interval = reconnection_interval + self.last_reconnection_time: None | datetime = None + self.disconnection_interval = 10.0 + self.notification_handler = notification_handler + self.client: None | Client = None + + SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + + self.metadata = { + 'model_name': '-', + 'model_id': '-', + 'workflow_name': 'opc_repository', + 'schedule_name': '-', + } + self._last_write_mono: float | None = None + self._connection_lock = asyncio.Lock() + self._session_ready = asyncio.Event() + self._reconnect_task: asyncio.Task[None] | None = None + self._allow_reconnect = True + + def _opc_debug_tags(self, session_id: str) -> dict[str, str]: + """ + Build Prometheus/log label tags for OPC session-scoped metrics. + + Args: + session_id (str): OPC UA session token string. + + Return: + dict[str, str]: Labels pod_id, server_name, runtime, opc_server_id, session_id. + """ + return { + 'pod_id': str(getattr(self, 'pod_id', 'unknown')), + 'server_name': self.server_name, + 'runtime': str(getattr(self, 'runtime', 'unknown')), + 'opc_server_id': self.id, + 'session_id': session_id, + } + + def _is_session_open(self) -> bool: + """ + Return whether the asyncua client has an open transport session. + + Return: + bool: True when protocol exists and is not closed. + """ + if self.client is None: + return False + try: + proto = self.client.uaclient.protocol + return proto is not None and proto.state != 'closed' + except Exception: + return False + + def _reconnection_window_elapsed(self) -> bool: + """ + Return whether enough time has passed since the last reconnect attempt. + + Return: + bool: True if a new reconnect is allowed. + """ + if self.last_reconnection_time is None: + return True + return ( + datetime.now() - self.last_reconnection_time + ).total_seconds() > self.reconnection_interval + + def _not_connected_error(self) -> dict[str, Any]: + """ + Build the standard error payload when validate_connection finds no open protocol. + + Return: + dict[str, Any]: Notification fields for OPC_CONNECTION_NOT_READY. + """ + return { + 'notification_id': f'OPC_CONNECTION_NOT_READY_{self.id}', + 'message': f'OPC server {self.id} is not connected', + 'block': 'opc_repository', + 'level': NotificationLevel.WARNING, + } + + async def set_security(self) -> None: + """ + Configure certificates and timeouts on the asyncua client. + + Raises: + ValueError: If cert paths or client are missing. + """ + if self.cert_path is None or self.private_key_path is None: + raise ValueError( + 'Certificate and private key paths must be provided for secure connection.' + ) + + cert = Path(self.cert_path) + private_key = Path(self.private_key_path) + server_cert = Path(self.server_cert_path) if self.server_cert_path else None + + if self.client is None: + raise ValueError('Client must be initialized before setting security') + + self.client.application_uri = self.server_uri + self.info('Setting security...', self.metadata) + await self.client.set_security( + SecurityPolicyBasic256, + certificate=str(cert), + private_key=str(private_key), + server_certificate=str(server_cert) if server_cert else None, + ) + self.client.secure_channel_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS + self.client.session_timeout = OPC_UA_SESSION_AND_CHANNEL_TIMEOUT_MS + + async def _create_client(self) -> None: + """ + Instantiate the asyncua Client and apply security when configured. + + Caller must hold _connection_lock. Does not open a UA session. + + Raises: + OpcClientAlreadyExistsError: If self.client is already set. + """ + if self.client is not None: + raise OpcClientAlreadyExistsError( + f'OPC client already exists for server {self.id}; ' + 'call disconnect() before creating a new client' + ) + + self.client = Client(self.url, timeout=10, watchdog_intervall=50) # type: ignore[attr-defined] + self.client.name = self.pod_id + self.client.application_name = self.pod_id + pod_uri = self.pod_id.replace('-', ':') + self.client.application_uri = pod_uri + self.client.product_uri = pod_uri + if self.cert_path: + await self.set_security() + + async def _open_session(self) -> tuple[bool, dict[str, Any]]: + """ + Open the OPC UA session on the existing client. + + Caller must hold _connection_lock. + + Raises: + OpcClientNotInitializedError: If self.client is None. + OpcSessionAlreadyConnectedError: If a session is already open. + + Return: + tuple[bool, dict[str, Any]]: Success flag and error payload on connect failure. + """ + if self.client is None: + raise OpcClientNotInitializedError( + f'OPC client is not initialized for server {self.id}; ' + 'call _create_client() before opening a session' + ) + if self._is_session_open(): + raise OpcSessionAlreadyConnectedError( + f'OPC session already connected for server {self.id}; ' + 'call disconnect() before connecting again' + ) + + tags = { + 'pod_id': self.pod_id, + 'server_name': self.server_name, + } + await self.emit_metric(metrics.OPC_CONNECTIONS_TOTAL, tags) + + try: + await self.client.connect() + + session_id = _opc_authentication_token_str(self.client) + revised_session_timeout_ms = int(self.client.session_timeout) + revised_secure_channel_timeout_ms = int(self.client.secure_channel_timeout) + self.info( + f'OPC new session connected opc_server_id={self.id} session_id={session_id} ' + f'revised_session_timeout_ms={revised_session_timeout_ms} ' + f'revised_secure_channel_timeout_ms={revised_secure_channel_timeout_ms}', + self.metadata, + ) + await self.emit_metric( + metrics.OPC_SESSION_CREATED_TOTAL, self._opc_debug_tags(session_id) + ) + await self.emit_metric( + metric_object=metrics.OPC_SESSION_REVISED_TIMEOUT_MS, + method='set', + tags=self._opc_debug_tags(session_id), + value=revised_session_timeout_ms, + ) + await self.emit_metric( + metric_object=metrics.OPC_CONNECTION_STATUS, + method='set', + tags={**tags, 'server_url': self.url}, + value=1, + ) + + self._last_write_mono = None + self._session_ready.set() + return True, {} + + except Exception as e: + await self._disconnect_locked() + trace = traceback.format_exc() + self.error(trace, self.metadata) + await self.emit_metric(metrics.OPC_CONNECTIONS_FAILED, tags) + return False, { + 'notification_id': f'OPC_CONNECTION_ERROR_{self.id}', + 'message': f'Failed to connect to OPC server: {e}', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + 'attachment_content': trace, + } + + async def _connect_locked(self) -> tuple[bool, dict[str, Any]]: + """ + Create the client when absent, then open a UA session. + + Caller must hold _connection_lock. + + Raises: + OpcSessionAlreadyConnectedError: If a session is already open. + + Return: + tuple[bool, dict[str, Any]]: Result from _open_session on connect failure. + """ + if self._is_session_open(): + raise OpcSessionAlreadyConnectedError( + f'OPC session already connected for server {self.id}; ' + 'call disconnect() before connecting again' + ) + if self.client is None: + await self._create_client() + return await self._open_session() + + async def _disconnection_fallback(self) -> list[dict[str, Any]]: + """ + Try up to five times to disconnect from the OPC UA server. + """ + assert self.client is not None + error_stack: list[dict[str, Any]] = [] + for i in range(5): + try: + self.info( + f'Disconnecting from OPC UA server, attempt {i + 1} of 5', + self.metadata, + ) + await self.client.disconnect() + return [] + except Exception as e: + self.error( + f'Failed to disconnect from OPC UA server in attempt {i + 1} of 5: {e}', + self.metadata, + ) + error_stack.append( + { + 'attempt': i + 1, + 'error': str(e), + 'traceback': traceback.format_exc(), + } + ) + await asyncio.sleep(self.disconnection_interval * i) + return error_stack + + async def _disconnect_locked(self) -> None: + """ + Tear down the current session and client. + + Caller must hold _connection_lock. + """ + self._last_write_mono = None + self._session_ready.clear() + + if self.client is None: + return + + session_id = _opc_authentication_token_str(self.client) + self.info( + f'OPC disconnecting opc_server_id={self.id} session_id={session_id}', + self.metadata, + ) + await self.emit_metric(metrics.OPC_SESSION_CLOSED_TOTAL, self._opc_debug_tags(session_id)) + + errors = await self._disconnection_fallback() + if errors: + await self.send_notification_async( + metadata=self.metadata, + notification_id=f'OPC_DISCONNECTION_ERROR_{self.id}', + message='Failed to disconnect from OPC server in 5 attempts.', + block='opc_repository', + level=NotificationLevel.ERROR, + attachment_content=json.dumps(errors, indent=4), + ) + else: + self.warning(f'Disconnected from OPC server {self.id} successfully', self.metadata) + + await self.emit_metric( + metric_object=metrics.OPC_CONNECTION_STATUS, + method='set', + tags={ + 'pod_id': self.pod_id, + 'server_name': self.server_name, + 'server_url': self.url, + }, + value=0, + ) + self.client = None + + async def _reconnect_locked(self) -> tuple[bool, dict[str, Any]]: + """ + Close the current session and open a new one. + + Caller must hold _connection_lock. Records last_reconnection_time for interval gating. + + Return: + tuple[bool, dict[str, Any]]: Result from _connect_locked after teardown. + """ + self.last_reconnection_time = datetime.now() + await self._disconnect_locked() + return await self._connect_locked() + + async def connect(self) -> tuple[bool, dict[str, Any]]: + """ + Open an OPC UA session under the connection lock (worker initialization). + """ + async with self._connection_lock: + self.info( + f'Starting connection to OPC server {self.id}:{self.server_name}...', + self.metadata, + ) + return await self._connect_locked() + + async def disconnect(self) -> None: + """ + Gracefully disconnect from the OPC server under the connection lock. + + Disables background reconnect so late writes during worker shutdown do not + respawn sessions. + """ + async with self._connection_lock: + self._allow_reconnect = False + await self._disconnect_locked() + + async def validate_connection(self) -> tuple[bool, dict[str, Any]]: + """ + Read-only check that the asyncua protocol is open. + + Caller must ensure _session_ready before writing. Does not connect or reconnect. + + Return: + tuple[bool, dict[str, Any]]: (True, {}) when open, otherwise (False, error). + """ + if self._is_session_open(): + return True, {} + self.error(f'OPC server {self.id} is not connected', self.metadata) + return False, self._not_connected_error() + + def _reconnect_task_in_progress(self) -> bool: + """ + Return whether a background reconnect task is currently running. + + Return: + bool: True when a reconnect task exists and has not finished. + """ + return self._reconnect_task is not None and not self._reconnect_task.done() + + async def _start_reconnect(self, reason: str, session_id: str) -> None: + """ + Schedule a background reconnect when allowed by interval and task state. + + Clears _session_ready before starting the task. No-op when _allow_reconnect is + False, the reconnection window has not elapsed, or a reconnect is already running. + + Args: + reason (str): Trigger for reconnect (OPC status name or synthetic reason). + session_id (str): Session token before failure. + """ + if not self._allow_reconnect: + return + if not self._reconnection_window_elapsed(): + self.warning( + f'OPC reconnect skipped reason=reconnection_window opc_server_id={self.id} ' + f'reconnect_reason={reason}', + self.metadata, + ) + return + if self._reconnect_task_in_progress(): + self.warning( + f'OPC reconnect skipped reason=in_progress opc_server_id={self.id} ' + f'reconnect_reason={reason}', + self.metadata, + ) + return + + self._session_ready.clear() + self.info( + f'OPC reconnect scheduled reconnect_reason={reason} opc_server_id={self.id} ' + f'old_session_id={session_id}', + self.metadata, + ) + self._reconnect_task = asyncio.create_task(self._run_reconnect(reason, session_id)) + + async def _run_reconnect(self, reason: str, session_id: str) -> None: + """ + Background task that tears down and re-establishes the OPC UA session. + + Args: + reason (str): Trigger for reconnect (OPC status or ProtocolClosed). + session_id (str): Previous session token string for logging. + """ + try: + async with self._connection_lock: + self.info( + f'OPC reconnect started reconnect_reason={reason} opc_server_id={self.id} ' + f'old_session_id={session_id}', + self.metadata, + ) + success, error = await self._reconnect_locked() + if not success: + self.error( + f'OPC reconnect failed reconnect_reason={reason} opc_server_id={self.id}', + self.metadata, + ) + if error: + self.error(error.get('message', ''), self.metadata) + except Exception: + self.error( + f'OPC reconnect task failed opc_server_id={self.id} reconnect_reason={reason}', + self.metadata, + ) + self.error(traceback.format_exc(), self.metadata) + + async def _log_write_inter_arrival(self, session_id: str, node: str) -> None: + """ + Log elapsed wall time since the previous successful OPC write on this repository. + + Args: + session_id (str): Current OPC UA session token string. + node (str): Node id written in this operation. + """ + now = time.monotonic() + if self._last_write_mono is not None: + delta_s = now - self._last_write_mono + self.info( + f'OPC write inter-arrival_s={delta_s:.6f} opc_server_id={self.id} ' + f'session_id={session_id} node={node}', + self.metadata, + ) + if self.client is not None: + session_timeout_ms = float(self.client.session_timeout) + if session_timeout_ms > 0 and delta_s > (session_timeout_ms / 1000.0): + await self.emit_metric( + metrics.OPC_WRITE_INTER_ARRIVAL_OVER_SESSION_TIMEOUT_TOTAL, + self._opc_debug_tags(session_id), + ) + self._last_write_mono = now + + async def _emit_opc_write_metric( + self, session_id: str, result: str, metadata: dict[str, Any] | None + ) -> None: + """ + Emit opc_write_attempts_total for a single write attempt outcome. + + Args: + session_id (str): OPC UA session token string, or "unknown". + result (str): Outcome label (OK, OPC status name, ProtocolClosed, etc.). + metadata (dict[str, Any] | None): Write context for model_id/model_name labels. + """ + await self.emit_metric( + metrics.OPC_WRITE_ATTEMPTS_TOTAL, + { + **self._opc_debug_tags(session_id), + **_model_labels_from_write_metadata(metadata), + 'result': result, + }, + ) + + def _write_failure_payload( + self, + notification_id: str, + message: str, + level: NotificationLevel = NotificationLevel.ERROR, + attachment_content: str | None = None, + opc_error_kind: str | None = None, + opc_status: str | None = None, + ) -> dict[str, Any]: + """ + Build a structured error dict returned from failed write_data paths. + + Args: + notification_id (str): Stable notification identifier. + message (str): Human-readable failure message. + level (NotificationLevel): Severity for downstream notifications. + attachment_content (str | None): Optional traceback or diagnostic text. + opc_error_kind (str | None): Classifier (session_bad, connection_lost, etc.). + opc_status (str | None): OPC UA status name or synthetic reason. + + Return: + dict[str, Any]: Error payload consumed by the OPC activity layer. + """ + payload: dict[str, Any] = { + 'notification_id': notification_id, + 'message': message, + 'block': 'opc_repository', + 'level': level, + } + if attachment_content is not None: + payload['attachment_content'] = attachment_content + if opc_error_kind is not None: + payload['opc_error_kind'] = opc_error_kind + if opc_status is not None: + payload['opc_status'] = opc_status + return payload + + async def _handle_tier1_bad( + self, + exc: BaseException, + session_id: str, + node: str, + metadata: dict[str, Any], + phase: str, + ) -> tuple[bool, dict[str, Any]]: + """ + Record metrics/logs and schedule reconnect after a Tier-1 Bad* error. + + Args: + exc (BaseException): Tier-1 OPC UA error. + session_id (str): Session token at failure time. + node (str): Node id being written. + metadata (dict[str, Any]): Write context. + phase (str): get_node or write_value. + + Return: + tuple[bool, dict[str, Any]]: Always (False, error payload). + """ + opc_status = _opc_status_from_exception(exc) + trace = traceback.format_exc() + self.error(trace, metadata) + await self._emit_opc_write_metric(session_id, opc_status, metadata) + self.error( + f'OPC write failed opc_status={opc_status} opc_server_id={self.id} ' + f'session_id={session_id} model_id={metadata.get("model_id", "unknown")} ' + f'model_name={metadata.get("model_name", "unknown")} node={node} phase={phase}', + metadata, + ) + await self._start_reconnect(opc_status, session_id) + return False, self._write_failure_payload( + notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}', + message=f'Failed to {phase} on OPC server: {exc} | metadata: {metadata}', + attachment_content=trace, + opc_error_kind='session_bad', + opc_status=opc_status, + ) + + async def _write_reconnect_in_progress( + self, metadata: dict[str, Any] + ) -> tuple[bool, dict[str, Any]]: + """ + Fail a write because a background reconnect task is already running. + + Args: + metadata (dict[str, Any]): Write context passed through to the activity. + + Return: + tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind reconnect_in_progress). + """ + await self._emit_opc_write_metric('unknown', 'ReconnectInProgress', metadata) + self.warning( + f'OPC write rejected reconnect_in_progress opc_server_id={self.id} ' + f'model_id={metadata.get("model_id", "unknown")} ' + f'model_name={metadata.get("model_name", "unknown")}', + metadata, + ) + return False, { + 'notification_id': f'OPC_WRITE_RECONNECT_IN_PROGRESS_{self.id}', + 'message': f'OPC write skipped: reconnect in progress | metadata: {metadata}', + 'block': 'opc_repository', + 'level': NotificationLevel.WARNING, + 'opc_error_kind': 'reconnect_in_progress', + } + + async def _write_connection_lost( + self, metadata: dict[str, Any], opc_status: str + ) -> tuple[bool, dict[str, Any]]: + """ + Fail a write after scheduling reconnect for a closed or stale session. + + Args: + metadata (dict[str, Any]): Write context passed through to the activity. + opc_status (str): Synthetic reason (ProtocolClosed, SessionNotReady). + + Return: + tuple[bool, dict[str, Any]]: (False, error info with opc_error_kind connection_lost). + """ + await self._emit_opc_write_metric('unknown', opc_status, metadata) + return False, self._write_failure_payload( + notification_id=f'OPC_WRITE_CONNECTION_LOST_{self.id}', + message=f'OPC write skipped: connection lost ({opc_status}) | metadata: {metadata}', + level=NotificationLevel.WARNING, + opc_error_kind='connection_lost', + opc_status=opc_status, + ) + + async def write_data( + self, node: str, value: Any, data_type: str, metadata: dict[str, Any] + ) -> tuple[bool, dict[str, Any]]: + """ + Write data to OPC server with a single attempt and background reconnect scheduling. + + Reconnect is scheduled on Tier-1 Bad*, closed protocol, or stale session readiness. + There is no retry within the same call. + + Args: + node (str): OPC UA node id to write. + value (Any): Value to convert and send. + data_type (str): Logical type key (float, int, bool, str, double). + metadata (dict[str, Any]): Activity context (model_id, model_name, etc.). + + Return: + tuple[bool, dict[str, Any]]: (True, {response_time}) on success, or + (False, structured error info) on failure. + """ + if self._reconnect_task_in_progress(): + return await self._write_reconnect_in_progress(metadata) + + if not self._session_ready.is_set(): + session_id = _opc_authentication_token_str(self.client) + await self._start_reconnect('SessionNotReady', session_id) + if self._reconnect_task_in_progress(): + return await self._write_reconnect_in_progress(metadata) + return await self._write_connection_lost(metadata, 'SessionNotReady') + + is_connected, _error = await self.validate_connection() + if not is_connected: + session_id = _opc_authentication_token_str(self.client) + await self._start_reconnect('ProtocolClosed', session_id) + return await self._write_connection_lost(metadata, 'ProtocolClosed') + + start_time = time.time() + session_id = _opc_authentication_token_str(self.client) + + try: + node_obj = self.client.get_node(node) # type: ignore[union-attr] + except Exception as e: + if is_reconnectable_opcua_bad(e): + return await self._handle_tier1_bad(e, session_id, node, metadata, 'get_node') + trace = traceback.format_exc() + self.error(trace, metadata) + await self._emit_opc_write_metric( + session_id, f'GetNodeError:{type(e).__name__}', metadata + ) + return False, self._write_failure_payload( + notification_id=f'OPC_WRITE_GET_NODE_ERROR_{self.id}', + message=f'Failed to get node from OPC server: {e} | metadata: {metadata}', + attachment_content=trace, + ) + + if data_type not in data_type_map: + await self._emit_opc_write_metric(session_id, 'UnsupportedDataType', metadata) + return False, self._write_failure_payload( + notification_id=f'OPC_WRITE_DATA_TYPE_ERROR_{self.id}', + message=f'Unsupported data type: {data_type} | metadata: {metadata}', + ) + + data = data_type_map[data_type]['converter'](value) + self.info(f'Writing {data} - {type(data)} to {node}', metadata) + ua_data = DataValue( + Variant(data, data_type_map[data_type]['opc_type']), + ) + + try: + await node_obj.write_value(ua_data) + end_time = time.time() + response_time = end_time - start_time + except Exception as e: + if is_reconnectable_opcua_bad(e): + return await self._handle_tier1_bad(e, session_id, node, metadata, 'write_value') + trace = traceback.format_exc() + self.error(trace, metadata) + await self._emit_opc_write_metric(session_id, type(e).__name__, metadata) + return False, self._write_failure_payload( + notification_id=f'OPC_WRITE_DATA_ERROR_{self.id}', + message=f'Failed to write data to OPC server: {e} | metadata: {metadata}', + attachment_content=trace, + ) + + await self._emit_opc_write_metric(session_id, 'OK', metadata) + await self._log_write_inter_arrival(session_id, node) + + return True, { + 'response_time': response_time, + } diff --git a/laborious/worker/__init__.py b/laborious/worker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py new file mode 100644 index 0000000..9f76141 --- /dev/null +++ b/laborious/worker/worker.py @@ -0,0 +1,336 @@ +""" +Laborious Worker Module + +This module provides the main worker implementation for the Sientia DataOps Laborious system. +It orchestrates Temporal workers, manages task queues, and handles the lifecycle of +prediction and retraining workflows. + +The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``): +- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload) +- minimal_retrain-{runtime}-queue: Model retraining workflows +- drift-{runtime}-queue: Drift detection workflows +- simple_metrics-{runtime}-queue: Simple metrics workflows +- import_model-{runtime}-queue: `.sientia` model import — registered **only** when + ``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision + +``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the +same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``). + +Key Features: +- Resource-based scaling with WorkerTuner (CPU and memory aware) +- Automatic polling scaling with PollerBehaviorAutoscaling +- Prometheus metrics integration +- Comprehensive error handling and logging +- Graceful shutdown with cleanup +- Multiple worker instances for different workflow types + +Environment Variables: +- RUNTIME: Required non-empty string; suffix for all task queue names +- TEMPORAL_HOST: Temporal server address (default: localhost:7233) +- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious) +- POD_ID: Kubernetes pod identifier for metrics +- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090) +- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091) +- PROJECT_NAME: Project name for notifications (default: laborious) +- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default +- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_*: import configuration, see + ``laborious.utils.connectors_config`` +""" + +from temporalio import client, workflow +from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig + +with workflow.unsafe.imports_passed_through(): + import asyncio + import os + import sys + + from prometheus_client import start_http_server + from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler + from sientia_do.observability.logger import get_logger + from sientia_do.temporal.worker.prepare_worker import prepare_worker + from sientia_do.utils.connectors_config import ( + build_api_config, + build_mongodb_config, + build_postgres_config, + ) + + from laborious import metrics + from laborious.activities.activities import Activities + from laborious.utils.connectors_config import ( + build_import_config, + build_import_status_config, + build_minio_config, + build_mlflow_config, + build_opc_config, + ) + from laborious.workflows.drift import Drift + from laborious.workflows.import_model import ImportModel + from laborious.workflows.minimal_retrain import MinimalRetrain + from laborious.workflows.predictions_batch import PredictionsBatch + from laborious.workflows.simple_metrics import SimpleMetrics + from laborious.workflows.sub_workflows.format_and_export_prediction import ( + FormatAndExportPrediction, + ) + from laborious.workflows.sub_workflows.prediction_process import PredictionProcess + +POD_ID = os.getenv('HOSTNAME') +SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) + + +async def main(): + """ + Main entry point for the Laborious worker application. + + This function initializes and starts all components of the worker: + 1. Sets up logging and metadata + 2. Starts Prometheus metrics server + 3. Initializes notification handler + 4. Creates and configures activities + 5. Initializes OPC connections + 6. Starts Temporal client and workers + 7. Manages worker lifecycle and graceful shutdown + + The function runs indefinitely until interrupted or an error occurs. + On error, it performs cleanup and exits with a non-zero status code. + + Raises: + Exception: Any unhandled exception during worker execution + SystemExit: On graceful shutdown or error conditions + """ + host = os.getenv('TEMPORAL_HOST', 'localhost:7233') + logger = get_logger(__name__) + + metadata = { + 'pod_id': POD_ID, + 'model_name': '-', + 'model_id': '-', + 'workflow_name': '-', + 'schedule_name': '-', + } + + logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata) + + runtime = os.getenv('RUNTIME', '').strip() + if not runtime: + logger.custom_critical( + 'RUNTIME environment variable is required and must be non-empty', + metadata, + ) + metrics.APP_UP.labels(pod_id=POD_ID).set(0) + sys.exit(1) + + metadata_runtime = {**metadata, 'runtime': runtime} + + logger.custom_info('Starting prometheus client...', metadata_runtime) + start_prometheus_server() + + logger.custom_info('Starting Notification Handler...', metadata_runtime) + + mongo_config = build_mongodb_config() + notification_handler = NotificationHandler( + connection_string=mongo_config['connection_string'], + database=mongo_config['database_name'], + logger=logger, + project_name=os.getenv('PROJECT_NAME', 'laborious'), + ) + + logger.custom_info('Starting Activities...', metadata_runtime) + + activities = Activities( + postgres_config=build_postgres_config(), + mlflow_config=build_mlflow_config(), + minio_config=build_minio_config(), + opc_config=build_opc_config(), + pi_web_api_config=build_api_config(), + logger=logger, + notification_handler=notification_handler, + mongo_config=mongo_config, + import_config=build_import_config(), + import_status_config=build_import_status_config(), + ) + + logger.custom_info('Initializing OPC...', metadata_runtime) + await activities.init_opc() + + logger.custom_info( + f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', + metadata_runtime, + ) + + new_runtime = Runtime( + telemetry=TelemetryConfig( + metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}') + ) + ) + + logger.custom_info(f'Starting Temporal Client at {host}...', metadata_runtime) + + temporal_client = await client.Client.connect( + target_host=host, + namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'), + runtime=new_runtime, + ) + + logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime) + + workers = [ + prepare_worker( + temporal_client=temporal_client, + main_workflow=MinimalRetrain, + other_workflows=[], + activities=[ + activities.load_query_with_minio_offload, + activities.retrain_model, + activities.update_production_model, + activities.format_retrain_report, + activities.export_data_to_postgres, + ], + logger=logger, + runtime=runtime, + ), + prepare_worker( + temporal_client=temporal_client, + main_workflow=SimpleMetrics, + other_workflows=[], + activities=[ + activities.load_custom_query, + activities.calculate_simple_metrics, + activities.export_data_to_postgres, + ], + logger=logger, + runtime=runtime, + ), + prepare_worker( + temporal_client=temporal_client, + main_workflow=Drift, + other_workflows=[], + activities=[ + activities.load_custom_query, + activities.get_reference_data, + activities.calculate_drift, + activities.export_data_to_postgres, + ], + logger=logger, + runtime=runtime, + ), + prepare_worker( + temporal_client=temporal_client, + main_workflow=PredictionsBatch, + other_workflows=[PredictionProcess, FormatAndExportPrediction], + activities=[ + # MLFlow + activities.request_predict, + activities.request_transform, + # Gates + activities.input_gate, + activities.mlflow_response_gate, + activities.mlflow_content_gate, + activities.format_transformed_data, + activities.format_prediction, + activities.format_default_prediction, + # OPC + activities.write_opc_data, + # Postgres / MinIO offload + activities.load_query_with_minio_offload, + activities.cleanup_minio_objects_expired, + activities.repeat_last_prediction, + activities.export_data_to_postgres, + activities.export_payload_to_postgres, + activities.write_metrics, + # Pi Web API + activities.write_pi_web_api_data, + ], + logger=logger, + runtime=runtime, + ), + ] + + # The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with + # its own credentials, and a flag lets that be a deployment decision instead of forcing every + # Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected + # and no import queue is polled at all. + if os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in ('1', 'true', 'yes', 'on'): + logger.custom_info('Import model worker enabled', metadata_runtime) + workers.append( + prepare_worker( + temporal_client=temporal_client, + main_workflow=ImportModel, + other_workflows=[], + activities=[ + # Import log row (the BFF database, update-only) + activities.claim_import_status, + activities.record_import_status, + activities.record_import_names, + activities.record_import_terminal_status, + activities.report_import_status_write_failure, + # Bundle + activities.download_import_bundle, + activities.open_import_bundle, + # MLflow provisioning + activities.create_import_experiment, + activities.upload_import_artifacts, + activities.register_import_model_version, + # The model listing + activities.write_import_model_document, + # Cleanup and retention + activities.cleanup_import_files, + activities.ensure_import_bundle_retention, + ], + logger=logger, + runtime=runtime, + ) + ) + + handlers = [] + for w in workers: + handlers.append(w.run()) + + logger.custom_info('Workers started successfully', metadata_runtime) + + exit_code = 0 + try: + # This will run the workers and wait for them to complete. + # If an exception occurs in any of the worker handlers, it will be propagated here. + await asyncio.gather(*handlers) + except BaseException as e: # NOSONAR + logger.custom_error(f'An unhandled exception occurred: {e}', metadata) + exit_code = 1 + finally: + if notification_handler: + notification_handler.shutdown() + if activities: + await activities.shutdown() + metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN + sys.exit(exit_code) + + +def start_prometheus_server(): + """ + Starts the Prometheus metrics server for monitoring and observability. + + This function initializes the Prometheus HTTP server on the configured port + and sets the application health metric to indicate the service is running. + + The server exposes metrics that can be scraped by Prometheus for monitoring + the health and performance of the Laborious worker. + + Environment Variables: + HTTP_METRICS_PORT: Port for the metrics server (default: 9090) + POD_ID: Pod identifier for metrics labeling + + Raises: + SystemExit: If the metrics server fails to start + """ + try: + port = int(os.getenv('HTTP_METRICS_PORT', 9090)) + start_http_server(port) + print(f'Prometheus server started on port {port}.') + metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP + except Exception as e: + print(f'Failed to start Prometheus server: {e}') + os._exit(1) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/laborious/workflows/__init__.py b/laborious/workflows/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py new file mode 100644 index 0000000..2faff32 --- /dev/null +++ b/laborious/workflows/drift.py @@ -0,0 +1,107 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + + +@workflow.defn(name='drift') +class Drift: + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the drift workflow. + + This method orchestrates the complete drift process by: + 1. Loading data using the provided custom SQL query + 2. Preparing prediction configuration and filters + 3. Delegating to the PredictionProcess workflow for ML operations + """ + + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'workflow_name': 'drift', + } + } + + print(f'Input data: {input_data}', metadata) + + model_config = input_data['model_config'] + target_name = model_config['target'] + + gathering_query = f""" + SELECT * + FROM "{input_data['schema']}"."{input_data['source_table_name']}" + WHERE + model_id = '{input_data['model_id']}' AND + timestamp > NOW() - INTERVAL '{input_data['interval']} minutes' + ORDER BY timestamp ASC + """ # nosec B608 - values come from internal Temporal workflow config, not user input + + target_data_handler = workflow.start_activity_method( + Activities.load_custom_query, + { + **metadata, + 'query': gathering_query, + 'datetime_columns': ['timestamp', 'created_at'], + 'orient': 'records', + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + reference_data_handler = workflow.start_activity_method( + Activities.get_reference_data, + {**metadata, 'model_name': input_data['model_name']}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + target_data = await target_data_handler + reference_data = await reference_data_handler + + if not target_data: + return + + drift_data = await workflow.execute_local_activity_method( + Activities.calculate_drift, + { + **metadata, + 'target_data': target_data, + 'reference_data': reference_data, + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'target_name': target_name, + 'drift_metrics': input_data.get( + 'drift_metrics', ['kolmogorov_smirnov', 'jensen_shannon', 'wasserstein'] + ), + 'chunk_period': input_data.get('chunk_period', 'min'), + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + if drift_data: + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'data': drift_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) diff --git a/laborious/workflows/import_model.py b/laborious/workflows/import_model.py new file mode 100644 index 0000000..39cbbf6 --- /dev/null +++ b/laborious/workflows/import_model.py @@ -0,0 +1,416 @@ +"""The `import_model` workflow: a `.sientia` object becomes a registered model, or a durable ERROR. + +Shape of this workflow, and the reasons it has that shape: + +- **The claim comes first.** The record already exists — the frontend inserted it through the BFF — + so the workflow's first act is to take ownership of it. Everything after that failure is recorded + *on* the record; a failed claim is recorded nowhere, because the workflow owns no row. +- **Fail-fast, and nothing else.** No rollback, no compensation, no resume, no reconciliation, no + `continue_as_new`. A failed import stays failed; trying again is a new import with a new upload. +- **Only the status writes retry**, and only because `UPDATE ... SET WHERE id` is + idempotent. Every provisioning call keeps `maximum_attempts=1`, because a repeated registry call or + Mongo insert can write twice. +- **One `finally`, two actions**: cleanup, then the single terminal write that classifies the + outcome once. No `except` block writes a verdict of its own. +- **The workflow holds ids, never contents.** The decrypted bundle never crosses an activity + boundary, and the password reaches this module only as an envelope it cannot open. +""" + +from datetime import timedelta + +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + from laborious.activities.model_import import ( + IMPORT_STATUS_HINT_RETRY_POLICY, + IMPORT_STATUS_RETRY_POLICY, + derive_run_name, + ) + from laborious.activities.model_import_errors import ImportInputError, failure_location + from laborious.utils.bundle.steps import ImportStep + +# `start_to_close_timeout` is per attempt. The terminal write additionally gets a +# `schedule_to_close_timeout` larger than its retry envelope (10 attempts, ≈151 s of backoff, ≈251 s +# if every attempt also burns its timeout) — set it smaller and Temporal cuts the retries short, +# which is the one configuration mistake that would silently undo design D20. +STATUS_WRITE_TIMEOUT = timedelta(seconds=10) +STATUS_WRITE_ENVELOPE = timedelta(seconds=300) +DOWNLOAD_TIMEOUT = timedelta(seconds=600) +OPEN_TIMEOUT = timedelta(seconds=900) +EXPERIMENT_TIMEOUT = timedelta(seconds=60) +ARTIFACT_UPLOAD_TIMEOUT = timedelta(seconds=900) +REGISTRATION_TIMEOUT = timedelta(seconds=120) +DOCUMENT_TIMEOUT = timedelta(seconds=60) +CLEANUP_TIMEOUT = timedelta(seconds=120) +RETENTION_TIMEOUT = timedelta(seconds=60) + +# Fields no import may carry. The project is a frontend concern: the model listing is the MongoDB +# `models` document, no migrated database holds a project or model table, and nothing in this +# repository reads a project-to-model link (design D22). An input naming one is rejected rather than +# accepted and ignored, so no caller believes a project link was recorded. The same goes for a +# relational target: there is none, and for a plaintext password, which does not belong in an event +# history whatever else is true. +REJECTED_INPUT_FIELDS = ( + 'project_id', + 'project_name', + 'provisioning_target', + 'schema', + 'table_name', + 'status_table', + 'password', +) + +DIGEST_LENGTH = 64 +_HEX_DIGITS = frozenset('0123456789abcdef') + + +@workflow.defn(name='import_model') +class ImportModel: + """Turn an encrypted `.sientia` object in MinIO into a registered model, or into an ERROR.""" + + @workflow.run + async def run(self, input_data: dict[str, Any]) -> dict[str, Any]: + """Execute the import. + + Args: + input_data: the import request. + Required keys: + - import_run_id (int): primary key of the import log row the frontend inserted + - object_key (str): the uploaded object's key + - expected_digest (str): the SHA-256 the uploader computed for those bytes + - password_envelope (str): the encrypted bundle password, never plaintext + Optional keys: + - bucket (str): defaults to the configured import bucket + + Returns: + dict: `import_run_id`, `run_id`, `model_name`, `version` and the model document's id. + + Raises: + ImportLogRowNotClaimableError: the record is not this workflow's — the run ends FAILED + with that error type and nothing is written to any row. + Exception: whatever failed. The record carries the failing step, its code and its + sentence; this exception carries the technical truth. + """ + workflow_id = workflow.info().workflow_id + import_run_id = self._require_import_run_id(input_data.get('import_run_id')) + metadata: dict[str, Any] = { + 'model_name': '-', + 'model_id': str(import_run_id), + 'workflow_name': 'import_model', + 'schedule_name': '-', + } + base = { + 'metadata': metadata, + 'import_run_id': import_run_id, + 'workflow_id': workflow_id, + } + + # Class A. Outside the try/finally on purpose: a claim that matches nothing leaves this + # workflow owning no row, so there is nothing to clean up and nothing to record. + claim = await workflow.execute_activity_method( + Activities.claim_import_status, + base, + retry_policy=IMPORT_STATUS_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + schedule_to_close_timeout=STATUS_WRITE_ENVELOPE, + ) + + current_step = ImportStep.RECEIVED + failed_step: str | None = None + gate: int | None = None + succeeded = False + source: dict[str, Any] = {} + result: dict[str, Any] = {'import_run_id': import_run_id} + + try: + self._validate_input(input_data, claim) + + current_step = ImportStep.DOWNLOAD + await self._hint(base, current_step) + await self._ensure_retention(base, input_data) + download = await workflow.execute_activity_method( + Activities.download_import_bundle, + { + **base, + 'bucket': input_data.get('bucket'), + 'object_key': input_data['object_key'], + 'expected_digest': input_data['expected_digest'], + }, + retry_policy=retry_policy, + start_to_close_timeout=DOWNLOAD_TIMEOUT, + ) + + current_step = ImportStep.DECRYPTION + await self._hint(base, current_step) + opened = await workflow.execute_activity_method( + Activities.open_import_bundle, + { + **base, + 'path': download['path'], + 'password_envelope': input_data['password_envelope'], + }, + retry_policy=retry_policy, + start_to_close_timeout=OPEN_TIMEOUT, + ) + source = opened['metadata'] + parameters = opened['parameters'] + model_name = source['model_name'] + metadata['model_name'] = model_name + run_name = derive_run_name(model_name, source['model_version']) + + # The names go on the record before anything exists in MLflow, so a person watching the + # import list sees *which* model is being imported while it is still running. + await self._record_names(base, source, run_name) + + current_step = ImportStep.EXPERIMENT_CREATION + await self._hint(base, current_step, source) + experiment = await workflow.execute_activity_method( + Activities.create_import_experiment, + {**base, 'experiment_name': source['experiment_name']}, + retry_policy=retry_policy, + start_to_close_timeout=EXPERIMENT_TIMEOUT, + ) + + current_step = ImportStep.ARTIFACT_UPLOAD + await self._hint(base, current_step, source) + run = await workflow.execute_activity_method( + Activities.upload_import_artifacts, + { + **base, + 'experiment_id': experiment['experiment_id'], + 'run_name': run_name, + 'extracted_dir': opened['extracted_dir'], + 'parameters': parameters, + 'source': source, + }, + retry_policy=retry_policy, + start_to_close_timeout=ARTIFACT_UPLOAD_TIMEOUT, + ) + + current_step = ImportStep.REGISTRATION + await self._hint(base, current_step, source) + version = await workflow.execute_activity_method( + Activities.register_import_model_version, + {**base, 'model_name': model_name, 'run_id': run['run_id']}, + retry_policy=retry_policy, + start_to_close_timeout=REGISTRATION_TIMEOUT, + ) + + current_step = ImportStep.MODEL_DOCUMENT + await self._hint(base, current_step, source) + document = await workflow.execute_activity_method( + Activities.write_import_model_document, + { + **base, + 'model_name': model_name, + 'target': parameters['target_variable'], + }, + retry_policy=retry_policy, + start_to_close_timeout=DOCUMENT_TIMEOUT, + ) + + succeeded = True + result = { + 'import_run_id': import_run_id, + 'run_id': run['run_id'], + 'model_name': model_name, + 'version': version['version'], + 'document_id': document.get('id'), + } + except BaseException as error: # NOSONAR - classified once below, then re-raised untouched + located_step, located_gate = failure_location(error) + failed_step = located_step or current_step.value + gate = located_gate + raise + finally: + cleanup = await self._cleanup(base) + await self._record_terminal( + base, + succeeded=succeeded, + step=failed_step or current_step.value, + gate=gate, + source=source, + cleanup_failed=not cleanup.get('cleaned', True), + ) + + return result + + # ------------------------------------------------------------------ validation + + @staticmethod + def _require_import_run_id(value: Any) -> int: + """The one check that runs before the claim, because the claim needs its result. + + A missing or unusable id means there is no record to claim and none to write: the run ends + FAILED with this error type, and the failure is visible in the workflow's terminal state, its + error type, the error metric and a notification — everywhere except a row, which is the + point. + """ + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ImportInputError('import_run_id must be a positive integer') + return value + + @staticmethod + def _validate_input(input_data: dict[str, Any], claim: dict[str, Any]) -> None: + """Validate the rest of the request, and check it against the record it just claimed. + + This runs **after** the claim so its failure lands on the record rather than only in the + logs: by now the row is this workflow's, so a bad request is reported to the person who made + it. Everything here fails at step `received`. + """ + unknown = [name for name in REJECTED_INPUT_FIELDS if name in input_data] + if unknown: + raise ImportInputError( + f'the import request carries fields this workflow does not accept: ' + f'{", ".join(sorted(unknown))}' + ) + + for required in ('object_key', 'expected_digest', 'password_envelope'): + if not input_data.get(required): + raise ImportInputError(f'the import request is missing {required}') + + digest = str(input_data['expected_digest']).lower() + if len(digest) != DIGEST_LENGTH or not set(digest) <= _HEX_DIGITS: + raise ImportInputError('expected_digest must be a 64-character hexadecimal SHA-256') + + object_key = input_data['object_key'] + file_name = claim.get('file_name') + if file_name and file_name != object_key: + raise ImportInputError('the record names a different object than the request') + + request_digest = claim.get('request_digest') + if request_digest and request_digest != digest: + raise ImportInputError('the record names a different digest than the request') + + # --------------------------------------------------------------- status writes + + async def _hint( + self, + base: dict[str, Any], + step: ImportStep, + source: dict[str, Any] | None = None, + ) -> None: + """Record the step about to be attempted. A hint: its failure never fails the import.""" + try: + await workflow.execute_activity_method( + Activities.record_import_status, + {**base, 'step': step.value, 'source': source or None}, + retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + ) + except BaseException as error: # NOSONAR - a lost hint must not fail a working import + await self._report_status_write_failure(base, error) + + async def _record_names( + self, base: dict[str, Any], source: dict[str, Any], run_name: str + ) -> None: + """Write the experiment and run names. Also a hint: the run is created either way.""" + try: + await workflow.execute_activity_method( + Activities.record_import_names, + { + **base, + 'experiment_name': source['experiment_name'], + 'run_name': run_name, + 'source': source, + 'step': ImportStep.EXPERIMENT_CREATION.value, + }, + retry_policy=IMPORT_STATUS_HINT_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + ) + except BaseException as error: # NOSONAR - the names are for a human reading the list + await self._report_status_write_failure(base, error) + + async def _record_terminal( + self, + base: dict[str, Any], + *, + succeeded: bool, + step: str | None, + gate: int | None, + source: dict[str, Any], + cleanup_failed: bool, + ) -> None: + """The single authoritative write, and the last thing this workflow does. + + Its own failure is reported through the channel with the verdict it could not write, and it + never turns a completed import into a failed one — nor masks the cause of a real failure. + """ + try: + await workflow.execute_activity_method( + Activities.record_import_terminal_status, + { + **base, + 'succeeded': succeeded, + 'step': step, + 'gate': gate, + 'source': source or None, + 'cleanup_failed': cleanup_failed, + }, + retry_policy=IMPORT_STATUS_RETRY_POLICY, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + schedule_to_close_timeout=STATUS_WRITE_ENVELOPE, + ) + except BaseException as error: # NOSONAR - the outcome stands even when it cannot be written + await self._report_status_write_failure( + base, + error, + verdict=f'succeeded={succeeded} step={step} gate={gate}', + ) + + async def _report_status_write_failure( + self, + base: dict[str, Any], + error: BaseException, + verdict: str | None = None, + ) -> None: + """Route a status write that could not be made through the one failure channel.""" + try: + await workflow.execute_activity_method( + Activities.report_import_status_write_failure, + {**base, 'detail': str(error), 'verdict': verdict}, + retry_policy=retry_policy, + start_to_close_timeout=STATUS_WRITE_TIMEOUT, + ) + except BaseException: # NOSONAR - reporting the report is where this stops + pass + + # ----------------------------------------------------------- side-effect steps + + @staticmethod + async def _ensure_retention(base: dict[str, Any], input_data: dict[str, Any]) -> None: + """Ensure the bundle bucket's expiry rule, without making it a reason to fail an import. + + The rule is applied by the importer rather than assumed from a Helm value, so the acceptance + criterion is verifiable from the code that depends on it. But a bucket policy the importer + cannot set is an operational problem, not a bad import: failing a valid model import because + a lifecycle API refused would be the wrong trade, so this is reported and dropped like a + lost status hint. + """ + try: + await workflow.execute_activity_method( + Activities.ensure_import_bundle_retention, + {**base, 'bucket': input_data.get('bucket')}, + retry_policy=retry_policy, + start_to_close_timeout=RETENTION_TIMEOUT, + ) + except BaseException: # NOSONAR - retention is not a precondition of a correct import + pass + + @staticmethod + async def _cleanup(base: dict[str, Any]) -> dict[str, Any]: + """Remove this import's temporary files. Runs on every path, before the terminal write.""" + try: + return await workflow.execute_activity_method( + Activities.cleanup_import_files, + base, + retry_policy=retry_policy, + start_to_close_timeout=CLEANUP_TIMEOUT, + ) + except BaseException: # NOSONAR - cleanup never replaces the failure that brought us here + return {'cleaned': False} diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py new file mode 100644 index 0000000..f8f5151 --- /dev/null +++ b/laborious/workflows/minimal_retrain.py @@ -0,0 +1,137 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + + +@workflow.defn(name='minimal_retrain') +class MinimalRetrain: + """ + Automated model retraining workflow for the Laborious system. + + This workflow implements a complete model retraining pipeline that loads + training data, executes model retraining, updates production models, + and maintains comprehensive audit trails. It's designed for automated + model lifecycle management with minimal manual intervention. + + The workflow provides a robust retraining process with: + - Automated data loading from configured data sources + - MLFlow model retraining with quality validation + - Production model updates with version control + - Comprehensive reporting and audit trail maintenance + - Error handling and notification integration + """ + + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the automated model retraining workflow. + + This method orchestrates the complete model retraining process by: + 1. Loading training data using the provided custom SQL query + 2. Executing MLFlow model retraining with the loaded data + 3. Updating production models with newly trained versions + 4. Persisting comprehensive retraining reports to database + + The method implements comprehensive error handling and ensures all + required parameters are properly configured before proceeding. + + Args: + input_data: Complete configuration for the retraining workflow + Required keys: + - schedule_name (str): Schedule identifier for the retraining + - model_name (str): Name of the ML model to retrain + - model_id (int): Unique identifier for the model version + - query (str): SQL query for training data loading + - schema (str, optional): Database schema for report storage + - table_name (str, optional): Target table for retraining reports + - datetime_columns (list[str], optional): Columns to treat as datetime + + Returns: + None: The workflow completes successfully when all steps finish + + Raises: + Exception: If any required parameters are missing or if the workflow fails + during data loading, retraining, or model update operations + """ + + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'workflow_name': 'minimal_retrain', + } + } + + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + + storage_result = await workflow.execute_activity_method( + Activities.load_query_with_minio_offload, + { + **metadata, + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': model_name, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=600), + ) + + storage_payload = MinioDataFramePayload.from_dict(storage_result) + if not storage_payload.has_data(): + raise ValueError('No data returned from query') + + experiment_response = await workflow.execute_activity_method( + Activities.retrain_model, + { + **metadata, + 'data': storage_result, + 'model_name': model_name, + 'model_config': model_config, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(hours=1), + ) + + if experiment_response['success']: + update_report = await workflow.execute_activity_method( + Activities.update_production_model, + {**metadata, 'model_name': model_name, **experiment_response}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + else: + update_report = {} + + report = await workflow.execute_local_activity_method( + Activities.format_retrain_report, + { + **metadata, + 'experiment_response': experiment_response, + 'model_name': model_name, + 'model_id': input_data['model_id'], + 'update_report': update_report, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'data': report, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=600), + ) diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py new file mode 100644 index 0000000..e68c215 --- /dev/null +++ b/laborious/workflows/predictions_batch.py @@ -0,0 +1,127 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + + +@workflow.defn(name='predictions_batch') +class PredictionsBatch: + """ + Main batch prediction workflow for the Laborious system. + + This workflow orchestrates the complete batch prediction process, handling + data loading, configuration management, and workflow delegation. It serves + as the primary entry point for batch prediction operations and ensures + proper data preparation before ML model inference. + + The workflow implements a robust data processing pipeline with: + - Custom SQL query execution for data loading + - Comprehensive configuration management + - Data quality filter application + - MLFlow model integration + - Workflow delegation to specialized sub-workflows + + Workflow Execution: + 1. Data Loading: Executes custom SQL query to load prediction data + 2. Configuration Preparation: Sets up prediction parameters and filters + 3. Workflow Delegation: Spawns PredictionProcess child workflow + 4. Error Handling: Implements comprehensive error handling and retry policies + """ + + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the batch prediction workflow. + + This method orchestrates the complete batch prediction process by: + 1. Loading data using the provided custom SQL query + 2. Preparing prediction configuration and filters + 3. Delegating to the PredictionProcess workflow for ML operations + + The method implements comprehensive error handling and ensures all + required parameters are properly configured before proceeding. + + Args: + input_data: Complete configuration for the batch prediction + Required keys: + - schedule_name (str): Schedule identifier for the prediction + - model_name (str): Name of the ML model to use + - model_id (int): Unique identifier for the model + - query (str): SQL query for data loading + - schema (dict, optional): Data schema definition + - table_name (str, optional): Target table for predictions + - input_filters (dict, optional): Data quality filters + - mlflow_transform_filters (dict, optional): MLFlow transform filters + - mlflow_predict_filters (dict, optional): MLFlow prediction filters + - model_retention (int, optional): Model retention period in minutes + - path_priority (list[str]): Decision path priority configuration + - opc_output_config (dict, optional): OPC server export configuration + - pi_web_api_output_config (dict, optional): PI Web API export configuration + - datetime_columns (list[str], optional): Columns to treat as datetime + - save_transform (bool, optional): Whether to save transformed data (default: True) + - prediction_store_policy (str, optional): Data retention policy (default: 'lts:1') + + Returns: + None: The workflow completes successfully when the child workflow finishes + + Raises: + Exception: If any required parameters are missing or if the workflow fails + during data loading or workflow delegation + """ + + metadata = { + 'metadata': { + 'schedule_name': input_data['schedule_name'], + 'model_name': input_data['model_name'], + 'model_id': input_data['model_id'], + 'workflow_name': 'predictions_batch', + } + } + + # Load data using custom query with optional MinIO offload for large frames + data = await workflow.execute_activity_method( + Activities.load_query_with_minio_offload, + { + **metadata, + 'query': input_data['query'], + 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + # Prepare input for prediction_process workflow + prediction_input = { + 'metadata': metadata, + 'data': data, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_table_name'], + 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], + 'input_filters': input_data.get( + 'input_filters', {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}} + ), + 'mlflow_transform_filters': input_data.get( + 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}} + ), + 'mlflow_predict_filters': input_data.get( + 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}} + ), + 'model_config': input_data.get('model_config', {}), + 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), + 'opc_output_config': input_data.get('opc_output_config', {}), + 'on_conflict': input_data.get('on_conflict', 'error'), + 'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}), + 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), + 'save_transform': input_data.get('save_transform', True), + } + + # Execute prediction process workflow + await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input) diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py new file mode 100644 index 0000000..38b47c4 --- /dev/null +++ b/laborious/workflows/simple_metrics.py @@ -0,0 +1,95 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + + +@workflow.defn(name='simple_metrics') +class SimpleMetrics: + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the simple metrics workflow. + """ + metadata = { + 'metadata': { + 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], + 'workflow_name': 'simple_metrics', + 'schedule_name': input_data['schedule_name'], + } + } + + model_id = input_data['model_id'] + interval_minutes = input_data['interval_minutes'] + + model_config = input_data['model_config'] + target_name = model_config['target'] + + query = f""" + select p."timestamp", p.prediction, ld.value as "target" + from "{input_data['schema']}"."{input_data['predictions_table_name']}" p + inner join "{input_data['schema']}"."{input_data['data_table_name']}" ld + on p."timestamp" = ld."timestamp" + where + p.model_id = '{model_id}' and + p.prediction is not null and + ld.variable = '{target_name}' and + ld.value is not null and + p."timestamp" >= NOW() - INTERVAL '{interval_minutes} minutes' + order by + p."timestamp" desc; + """ # nosec B608 - values come from internal Temporal workflow config, not user input + + target_data = await workflow.execute_activity_method( + Activities.load_custom_query, + { + **metadata, + 'query': query, + 'datetime_columns': ['timestamp'], + 'orient': 'records', + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + if not target_data: + return + + simple_metrics = await workflow.execute_local_activity_method( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': model_id, + 'target_data': target_data, + 'metrics': input_data.get('metrics', ['rmse', 'mse', 'mae', 'r2']), + 'interval_minutes': interval_minutes, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) + + if not simple_metrics: + return + + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'data': simple_metrics, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=300), + ) diff --git a/laborious/workflows/sub_workflows/__init__.py b/laborious/workflows/sub_workflows/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py new file mode 100644 index 0000000..68552b2 --- /dev/null +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -0,0 +1,220 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + + +@workflow.defn(name='subworkflow.format_and_export_prediction') +class FormatAndExportPrediction: + """ + Data formatting and export workflow for prediction results. + + This workflow handles the final stages of the prediction pipeline, including + data formatting, database persistence, OPC server export, and metrics recording. + It implements flexible formatting based on prediction quality and provides + comprehensive export capabilities to multiple destinations. + + The workflow supports two main prediction paths: + 1. Normal Prediction: Formats and exports successful prediction results + 2. Default Prediction: Creates fallback predictions for error conditions + + Export Destinations: + - PostgreSQL Database: Persistent storage with timestamp conversion + - PI Web API: Real-time industrial system integration for prediction and confidence values + - OPC Servers: Real-time industrial system integration + - Prometheus Metrics: Performance monitoring and operational visibility + """ + + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the prediction formatting and export workflow. + + This method orchestrates the complete data export process by: + 1. Determining the appropriate formatting strategy based on path_flag + 2. Formatting prediction data according to quality and requirements + 3. Exporting data to PI Web API for real-time industrial access (if configured) + 4. Exporting data to OPC servers for real-time industrial access (if configured) + 5. Persisting data to PostgreSQL database with comprehensive metadata + 6. Recording performance metrics for operational monitoring + + The method implements flexible formatting strategies: + - Normal predictions: Full data formatting with confidence scores + - Error predictions: Default formatting with error indicators + - Comprehensive export: Multi-destination data distribution + + Args: + input_data: Complete configuration for the export workflow + Required keys: + - metadata (dict): Workflow execution metadata + - path_flag (str | None): Decision path flag for formatting strategy + - None: Normal prediction path with full formatting + - Any other value: Default prediction path for error conditions + - data (dict[str, Any]): Prediction data to format and export + - prediction_confidence (float): Confidence score for the prediction + - timestamp (str): ISO-formatted timestamp for the prediction + - model_id (int): Unique identifier for the ML model + - model_name (str): Name of the ML model + - schema (str): Database schema for data storage + - table_name (str): Target table for data persistence + Optional keys: + - opc_output_config (dict[str, Any]): OPC server export configuration + - pi_web_api_output_config (dict[str, Any]): PI Web API export configuration + Contains endpoint, prediction_tags, and confidence_tags mappings + - transformed_data (dict[str, Any]): Transformed data to export separately + Only processed when path_flag is None + - transform_table_name (str): Target table for transformed data export + Required if transformed_data is provided + - prediction_store_policy (str): Data retention policy (e.g., 'lts:1', 'erl:2') + Required when path_flag is None + - comment (str): Operational comment or error description + Required when path_flag is not None + + Returns: + None: The workflow completes successfully when all export operations finish + + Note: + When transformed_data is provided and path_flag is None, the workflow will: + 1. Format the transformed data using format_transformed_data + 2. Export it to a separate table (transform_table_name) asynchronously + 3. Wait for both prediction and transformed data exports to complete + """ + metadata = input_data['metadata'] + path_flag = input_data['path_flag'] + data = input_data['data'] + transformed_data = input_data.get('transformed_data', None) + prediction_confidence = input_data['prediction_confidence'] + + opc_output_config = input_data.get('opc_output_config', None) + pi_web_api_output_config = input_data.get('pi_web_api_output_config', None) + + if path_flag is None: + # Normal prediction path: format prediction data with full metadata + prediction = await workflow.execute_local_activity_method( + Activities.format_prediction, + { + **metadata, + 'data': data, + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': prediction_confidence, + 'prediction_store_policy': input_data['prediction_store_policy'], + 'model_name': input_data['model_name'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + # Optionally format and export transformed data to separate table + if transformed_data is not None: + transformed = await workflow.execute_local_activity_method( + Activities.format_transformed_data, + { + **metadata, + 'data': transformed_data, + 'model_id': input_data['model_id'], + 'model_name': input_data['model_name'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + write_transformed_handler = workflow.start_activity_method( + Activities.export_payload_to_postgres, + { + **metadata, + 'schema': input_data['schema'], + 'table_name': input_data['transform_table_name'], + 'data': transformed, + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + else: + write_transformed_handler = None + + else: + # Error path: create default prediction with error indicators + prediction = await workflow.execute_local_activity_method( + Activities.format_default_prediction, + { + **metadata, + 'timestamp': input_data['timestamp'], + 'model_id': input_data['model_id'], + 'prediction_confidence': prediction_confidence, + 'comment': input_data['comment'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + write_transformed_handler = None + + opc_metrics = {} + + # write to pi web api + if pi_web_api_output_config: + prediction = await workflow.execute_activity_method( + Activities.write_pi_web_api_data, + { + 'pi_web_api_output_config': pi_web_api_output_config, + 'data': prediction, + **metadata, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + # write to opc + if opc_output_config: + prediction, opc_metrics = await workflow.execute_activity_method( + Activities.write_opc_data, + { + 'opc_output_config': opc_output_config, + 'data': prediction, + **metadata, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) + + # write to postgres + await workflow.execute_activity_method( + Activities.export_data_to_postgres, + { + **metadata, + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'data': prediction, + 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + 'on_conflict': input_data.get('on_conflict', 'error'), + 'unique_columns': ['model_id', 'timestamp'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=180), + ) + + if write_transformed_handler is not None: + await write_transformed_handler + + await workflow.execute_activity_method( + Activities.write_metrics, + { + **metadata, + 'prediction': prediction, + 'opc_metrics': opc_metrics, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=60), + ) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py new file mode 100644 index 0000000..da4125b --- /dev/null +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -0,0 +1,346 @@ +from temporalio import workflow + +with workflow.unsafe.imports_passed_through(): + from datetime import timedelta + from typing import Any + + from sientia_do.temporal.policies import retry_policy + + from laborious.activities.activities import Activities + + +@workflow.defn(name='subworkflow.prediction_process') +class PredictionProcess: + """ + Core prediction processing workflow for the Laborious system. + + This workflow implements the complete ML model inference pipeline, handling + data quality validation, MLFlow model interactions, and prediction processing. + It serves as the central orchestrator for all prediction operations and ensures + data quality throughout the entire process. + + The workflow implements a robust data processing pipeline with: + - Data quality validation using configurable filters + - MLFlow model transformation and prediction + - Response validation and quality assurance + - Flexible decision path handling + - Comprehensive error handling and retry policies + + Workflow Execution: + 1. Timestamp Retrieval: Gets last processed timestamp for incremental processing + 2. Input Data Gate: Applies data quality filters + 3. Path Decision: Determines processing path based on filter results + 4. MLFlow Transform: Requests data transformation using MLFlow models + 5. Response Validation: Filters transform responses for quality assurance + 6. MLFlow Prediction: Executes prediction using transformed data + 7. Content Validation: Filters prediction responses for final quality check + 8. Export Delegation: Delegates to FormatAndExportPrediction workflow + """ + + @workflow.run + async def run(self, input_data: dict[str, Any]): + """ + Execute the prediction process workflow. + + This method orchestrates the complete prediction processing pipeline by: + 1. Retrieving the last processed timestamp for incremental processing + 2. Applying data quality filters to validate input data + 3. Executing MLFlow model transformation and prediction + 4. Validating all responses for quality assurance + 5. Delegating to export workflow for data persistence + + The method implements comprehensive error handling and ensures all + data quality requirements are met before proceeding with ML operations. + + Args: + input_data: Complete configuration for the prediction process + Required keys: + - metadata (dict): Workflow execution metadata + - data (dict): Input data for prediction processing + - schema (dict): Data schema definition + - table_name (str): Target table for predictions + - model_id (str): ML model identifier + - model_name (str): ML model name + - input_filters (dict): Data quality filters + - mlflow_transform_filters (dict): MLFlow transform filters + - mlflow_predict_filters (dict): MLFlow prediction filters + - model_retention (int): Model retention period in minutes + - path_priority (list[str]): Decision path priority configuration + - opc_output_config (dict, optional): OPC server export configuration + - pi_web_api_output_config (dict, optional): PI Web API export configuration + - save_transform (bool, optional): Whether to save transformed data (default: True) + - prediction_store_policy (str, optional): Data retention policy (default: 'lts:1') + + Returns: + None: The workflow completes successfully when export workflow finishes + + Raises: + Exception: If any required parameters are missing or if the workflow fails + during data processing, MLFlow operations, or workflow delegation + + """ + + metadata = input_data['metadata'] + data = input_data['data'] + model_id = input_data['model_id'] + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + save_transform = input_data.get('save_transform', True) + + try: + await self._run_prediction_pipeline( + input_data, + metadata, + data, + model_id, + model_name, + model_config, + save_transform, + ) + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + except Exception as e: + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + raise e + + async def _run_prediction_pipeline( + self, + input_data: dict[str, Any], + metadata: dict[str, Any], + data: dict[str, Any], + model_id: str, + model_name: str, + model_config: dict[str, Any], + save_transform: bool, + ) -> None: + last_timestamp = data['last_timestamp'] + + # Apply input data quality gates + gate_input = { + **metadata, + 'filters': input_data['input_filters'], + 'data': data, + 'path_priority': input_data['path_priority'], + } + + path_flag, confidence, comment = await workflow.execute_local_activity_method( + Activities.input_gate, + gate_input, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=1), + ) + + # Handle path decision based on filter results + if await self.path_flag_handler( + data, path_flag, input_data, confidence, last_timestamp, comment + ): + return + + # Request MLFlow model transformation + transformed_data = await workflow.execute_activity_method( + Activities.request_transform, + {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + + # Validate MLFlow transform response + path_flag, confidence, comment = await workflow.execute_local_activity_method( + Activities.mlflow_response_gate, + { + **metadata, + 'filters': input_data['mlflow_transform_filters'], + 'data': transformed_data, + 'type': 'transform', + 'path_priority': input_data['path_priority'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=1), + ) + + # Handle path decision based on transform validation + if await self.path_flag_handler( + data, path_flag, input_data, confidence, last_timestamp, comment + ): + return + + path_flag, confidence, comment = await workflow.execute_local_activity_method( + Activities.mlflow_content_gate, + { + **metadata, + 'filters': input_data['mlflow_transform_filters'], + 'data': transformed_data, + 'type': 'transform', + 'path_priority': input_data['path_priority'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=1), + ) + + if await self.path_flag_handler( + data, path_flag, input_data, confidence, last_timestamp, comment + ): + return + + predicted_data = await workflow.execute_activity_method( + Activities.request_predict, + { + **metadata, + 'data': transformed_data, + 'model_name': model_name, + 'model_config': model_config, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + + # Validate MLFlow prediction response + path_flag, confidence, comment = await workflow.execute_local_activity_method( + Activities.mlflow_response_gate, + { + **metadata, + 'filters': input_data['mlflow_predict_filters'], + 'data': predicted_data, + 'type': 'predict', + 'path_priority': input_data['path_priority'], + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=1), + ) + + # Handle path decision based on prediction validation + if await self.path_flag_handler( + data, path_flag, input_data, confidence, last_timestamp, comment + ): + return + + # Delegate to export workflow for data persistence + await workflow.execute_child_workflow( + 'subworkflow.format_and_export_prediction', + { + 'metadata': metadata, + 'on_conflict': input_data.get('on_conflict', 'error'), + 'path_flag': path_flag, + 'data': predicted_data, + 'transformed_data': transformed_data if save_transform else None, + 'prediction_confidence': confidence, + 'timestamp': last_timestamp, + 'model_id': model_id, + 'model_name': model_name, + 'model_config': model_config, + 'opc_output_config': input_data['opc_output_config'], + 'pi_web_api_output_config': input_data['pi_web_api_output_config'], + 'schema': input_data['schema'], + 'table_name': input_data['table_name'], + 'transform_table_name': input_data['transform_table_name'], + 'comment': comment, + 'prediction_store_policy': input_data['prediction_store_policy'], + }, + ) + + async def path_flag_handler( + self, + data: dict[str, Any], + path_flag: str, + input_data: dict, + confidence: int, + last_timestamp: str, + comment: str, + ) -> bool: + """ + Handle path decisions based on filter results and confidence levels. + + This method determines the appropriate action based on the path flag + returned by data quality filters. It can stop processing, continue, + or repeat operations based on the configured path priority. + + Args: + data: Input data for processing + path_flag: Path decision from filter (STOP, CONTINUE, REPEAT) + input_data: Complete workflow input configuration including: + - metadata (dict): Workflow execution metadata + - schema (str): Database schema + - table_name (str): Target table for predictions + - transform_table_name (str): Target table for transformed data + - model_id (str): ML model identifier + - model_name (str): ML model name + - model_config (dict, optional): Model configuration + - opc_output_config (dict, optional): OPC server export configuration + - pi_web_api_output_config (dict, optional): PI Web API export configuration + - prediction_store_policy (str, optional): Data retention policy + confidence: Confidence level from filter validation + last_timestamp: Last processed timestamp + comment: Additional information about the filter result + + Returns: + bool: True if processing should stop, False to continue + + Path Handling: + - STOP: Terminates workflow execution + - CONTINUE: Delegates to FormatAndExportPrediction workflow with current data + - REPEAT: Repeats last prediction if available + """ + metadata = input_data['metadata'] + + schema = input_data['schema'] + table_name = input_data['table_name'] + transform_table_name = input_data['transform_table_name'] + model_id = input_data['model_id'] + model_name = input_data['model_name'] + model_config = input_data.get('model_config', {}) + + path_flag = path_flag.upper() if path_flag else '' + + if path_flag == 'STOP': + # Stop processing and exit workflow + return True + elif path_flag == 'REPEAT': + # Repeat last prediction if available + await workflow.execute_activity_method( + Activities.repeat_last_prediction, + { + **metadata, + 'schema': schema, + 'table_name': table_name, + 'model': model_id, + 'last_timestamp': last_timestamp, + }, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=1), + ) + return True + elif path_flag == 'CONTINUE': + # call write workflow + await workflow.execute_child_workflow( + 'subworkflow.format_and_export_prediction', + { + 'metadata': metadata, + 'path_flag': path_flag, + 'data': data, + 'prediction_confidence': confidence, + 'timestamp': last_timestamp, + 'model_id': model_id, + 'model_name': model_name, + 'model_config': model_config, + 'schema': schema, + 'table_name': table_name, + 'transform_table_name': transform_table_name, + 'comment': comment, + 'opc_output_config': input_data['opc_output_config'], + 'pi_web_api_output_config': input_data['pi_web_api_output_config'], + 'prediction_store_policy': input_data['prediction_store_policy'], + 'on_conflict': input_data.get('on_conflict', 'error'), + }, + ) + return True + + return False