From 981ac700d434407cc708bef6d0485d7b15d3ef2b Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 19 Mar 2026 17:29:43 -0300 Subject: [PATCH 01/51] SIENTIAPDE-1712 Implement MinIO Offload and Retention Features - Added configuration options for MinIO retention hours and offload threshold in README. - Introduced MinIO payload offloading for large DataFrame-derived payloads, storing them as parquet files. - Updated activities to utilize MinIO for data loading and cleanup, including new methods for offloading and retention management. - Refactored existing activities to integrate MinIO functionality, ensuring compatibility with previous workflows. - Removed the legacy MinioRepository class, consolidating MinIO operations under a new manager structure. - Updated requirements to use the latest version of the sientia-dataops-library. --- README.md | 70 +++++ e2e/conftest.py | 30 ++- laborious/activities/activities.py | 18 +- laborious/activities/gates.py | 105 ++++---- laborious/activities/mlflow.py | 118 ++++++--- laborious/activities/storage.py | 216 ++++++++++++--- laborious/utils/connectors_config.py | 1 + laborious/utils/models/__init__.py | 0 .../utils/models/minio_dataframe_payload.py | 230 ++++++++++++++++ laborious/utils/repository/minio_manager.py | 26 ++ .../utils/repository/minio_repository.py | 215 --------------- .../utils/repository/model_repository.py | 6 +- laborious/worker/worker.py | 6 +- laborious/workflows/minimal_retrain.py | 7 +- laborious/workflows/predictions_batch.py | 7 +- .../format_and_export_prediction.py | 4 +- .../sub_workflows/prediction_process.py | 53 +++- requirements.txt | 2 +- tests/laborious/activities/test_mlflow.py | 103 +++++--- tests/laborious/activities/test_storage.py | 122 +++++++-- .../models/test_minio_dataframe_payload.py | 54 ++++ .../utils/repository/test_minio_repository.py | 245 ------------------ .../subworkflows/test_prediction_process.py | 5 + .../workflows/test_minimal_retrain.py | 19 +- .../workflows/test_predictions_batch.py | 13 +- 25 files changed, 994 insertions(+), 681 deletions(-) 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 delete mode 100644 laborious/utils/repository/minio_repository.py create mode 100644 tests/laborious/utils/models/test_minio_dataframe_payload.py delete mode 100644 tests/laborious/utils/repository/test_minio_repository.py diff --git a/README.md b/README.md index 3041bfc..a02f756 100644 --- a/README.md +++ b/README.md @@ -716,10 +716,80 @@ The Laborious system exposes comprehensive Prometheus metrics for operational vi | `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No | | `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No | | `POD_ID` | Kubernetes pod identifier | `None` | No | +| `SIENTIA_MINIO_RETENTION_HOURS` | Retention window for offloaded MinIO objects | `168` | No | +| `SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES` | Offload threshold for DataFrame-derived payloads | `int(1.5 * 1024 * 1024)` | No | +### MinIO Payload Offload & Retention + +Laborious uses MinIO to prevent Temporal workflow history from carrying very large in-memory payloads (pandas `DataFrame`-derived dicts). +Whenever a payload exceeds a configurable size threshold, it is stored as a parquet file in MinIO and the workflow history only keeps a lightweight reference. + +Notes: +- `SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES` supports: + - Integer bytes (e.g. `"1572864"`) + - Float MiB (e.g. `"1.5"`), converted to bytes as `MiB * 1024 * 1024` +- Fallback behavior uses `1.5 MiB` when the env var is missing or invalid. + +#### Wire Contract: `MinioDataFramePayload` + +The payload is implemented in `laborious/utils/models/minio_dataframe_payload.py`. +The dataclass does **not** store a pandas `DataFrame` field. +Instead, the `DataFrame` is only used at build time by: +- `MinioDataFramePayload.from_dataframe(...)` +- `MinioDataFramePayload.from_dataframe_to_dict(...)` + +After evaluation, the payload is serialized for Temporal as a flat dict: +- **Inline path**: `data` contains `df.to_dict()`, and MinIO keys (`object_key`, `bucket`, ...) are absent / `None`. +- **MinIO path**: the dict contains: + - `bucket` + - `object_key` (full MinIO object name returned by `MinioRepository.upload_file`) + - `object_prefix` (directory prefix used for cleanup listing; relative to the repository namespace) + - `uri` (best-effort `s3:///<...>` string) + - `data` is omitted / set to `None`. + +When an activity needs pandas operations, it resolves references using: +- `MinioDataFramePayload.dataframe_from_wire(...)` + +#### MinIO Object Naming (Retention Parsing) + +MinIO object basename (required convention): +`{model_name}-{operation}-{timestamp}.parquet` + +Where: +- `model_name`: model identifier used by the pipeline +- `operation`: `initial` (SQL/query load before transform) or `transform` (after MLFlow transform) +- `timestamp`: `DATETIME_FORMAT_FILENAME` from `sientia_do.temporal.constants` + +The relative object key (under the repository namespace) is always shaped as: +`training_datasets/{model_name}/{basename}` + +Retention cleanup parses timestamps from the basename using the `-initial-` / `-transform-` anchors. +`model_name` may contain hyphens; parsing is resilient to it. + +#### Workflows / Activities Integration + +Predictions batch uses MinIO offload as follows: +1. `predictions_batch` calls `Activities.load_query_with_minio_offload` + - On success, it puts the serialized `MinioDataFramePayload` dict into `prediction_input["data"]`. +2. `sub_workflows/prediction_process` + - Tracks which MinIO prefixes were referenced for offloaded payloads. + - Runs `Activities.cleanup_minio_objects_expired` in a `finally` block (only when MinIO offload happened). +3. `laborious/activities/gates.py` and `laborious/activities/mlflow.py` + - Resolve offloaded payloads transparently before constructing pandas `DataFrame` objects. + +#### Legacy: `query_to_minio` (Minimal Retrain) + +`Storage.query_to_minio` is intentionally kept with its legacy behavior for `minimal_retrain`. +It always uploads parquet and returns `{success, object_key, uri}`. +It is not used by predictions batch MinIO offload, and its objects are not part of the retention parser described above. + +Legacy MinIO object layout (relative key): +`training_datasets/{model_name}/{object_prefix}_{timestamp}.parquet` where `object_prefix` is sanitized +(slashes replaced by underscores) to keep a stable model-level directory. + ### OPC Configuration For multiple OPC servers, use the `OPC_CONFIG` environment variable: diff --git a/e2e/conftest.py b/e2e/conftest.py index 5dd208e..ca3830e 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -3,6 +3,7 @@ Pytest configuration and fixtures for E2E tests. """ from unittest.mock import AsyncMock, MagicMock, patch +from io import BytesIO import pandas as pd import pytest @@ -216,11 +217,28 @@ def metrics_controller(mock_logger): def mock_minio_repository(): """Mock MinIO repository for object storage operations.""" mock_repo = MagicMock() + + # Provide at least valid parquet bytes so that MinioDataFramePayload.retrieve() + # can decode the payload if offloading is exercised in an integration scenario. + parquet_df = pd.DataFrame({'a': [1]}) + parquet_buffer = BytesIO() + parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True) + parquet_bytes = parquet_buffer.getvalue() - # Mock repository methods - mock_repo.put_parquet_from_dataframe = AsyncMock(return_value='test-object-key') - mock_repo.get_parquet_as_dataframe = AsyncMock(return_value=pd.DataFrame()) - mock_repo.minio_bucket = 'test-bucket' + # sientia_do MinioRepository API + mock_repo.bucket = 'test-bucket' + mock_repo.upload_file = AsyncMock( + side_effect=lambda file_bytes, relative_key, content_type='application/octet-stream', bucket=None, metadata=None: { + 'minio_object_name': f'sientia/streamlit-connectors/{relative_key}', + 'original_filename': relative_key.rsplit('/', 1)[-1], + 'uploaded_at': '2024-01-01T00:00:00Z', + 'sha256_hash': 'deadbeef', + } + ) + mock_repo.download_file = AsyncMock(return_value=parquet_bytes) + mock_repo.list_objects = AsyncMock(return_value=[]) + mock_repo.delete_file = AsyncMock() + mock_repo.close = MagicMock() return mock_repo @@ -261,7 +279,7 @@ def patch_create_engine(postgres_engine): @pytest_asyncio.fixture def patch_minio_repository(mock_minio_repository): """Patch MinioRepository to return mock.""" - with patch('laborious.utils.repository.minio_repository.MinioRepository', return_value=mock_minio_repository): + with patch('sientia_do.repository.minio_repository.MinioRepository', return_value=mock_minio_repository): yield @pytest_asyncio.fixture @@ -433,6 +451,8 @@ async def temporal_worker(temporal_test_env, test_activities): workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], activities=[ test_activities.load_custom_query, + test_activities.load_query_with_minio_offload, + test_activities.cleanup_minio_objects_expired, test_activities.get_last_timestamp, test_activities.input_gate, test_activities.request_transform, diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index f009b10..f714a4f 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -6,6 +6,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger + from sientia_do.repository.minio_repository import MinioRepository from laborious.activities.api import API from laborious.activities.gates import Gates @@ -13,6 +14,7 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.model_metrics import ModelMetrics from laborious.activities.opc import OPC from laborious.activities.storage import Storage + class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): @@ -73,6 +75,16 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): """ metrics_controller = MetricsController(logger=logger) + minio_repository = MinioRepository( + endpoint_url=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, + ) + # Initialize parent classes Storage.__init__( self, @@ -83,7 +95,8 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): dbname=postgres_config['dbname'], min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], - minio_config=minio_config, + retention_hours=minio_config['retention_hours'], + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -95,7 +108,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): mlflow_port=mlflow_config['port'], mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], - minio_config=minio_config, + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -103,6 +116,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): Gates.__init__( self, + minio_repository=minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index b581920..f226cc8 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -1,5 +1,8 @@ +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 @@ -15,6 +18,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.utils.formatters import create_sample_dict from laborious import metrics + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload from laborious.utils.filters.conditional_filters import ( filter_empty_data, filter_specific_variables_null_values, @@ -63,7 +67,7 @@ mlflow_content_path_confidence: Mapping[str, int] = { } -class Gates(SientiaMonitoring): +class Gates(MinioManager): """ Data quality gates and filtering activities for the Laborious system. @@ -83,11 +87,14 @@ class Gates(SientiaMonitoring): mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions """ + minio_repository: MinioRepository | None = None + def __init__( self, - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, + 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. @@ -99,13 +106,14 @@ class Gates(SientiaMonitoring): Raises: Exception: If BaseActivity initialization fails """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + MinioManager.__init__(self, minio_repository, logger, notification_handler, metrics_controller) def close(self) -> None: """ Close the gates activity and clean up resources. """ - SientiaMonitoring.shutdown(self) + + MinioManager.close(self) def __del__(self): self.close() @@ -149,7 +157,8 @@ class Gates(SientiaMonitoring): self.info('Performing input gate...', metadata) filters = input_data['filters'] - data = DataFrame(input_data['data']) + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) path_priority = input_data['path_priority'] filter_output = [] @@ -183,6 +192,9 @@ class Gates(SientiaMonitoring): 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') @@ -223,7 +235,10 @@ class Gates(SientiaMonitoring): self.info('Performing mlflow response gate...', metadata) filters = input_data['filters'] - data = input_data['data'] + + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) + gate_type = input_data['type'] path_priority = input_data['path_priority'] @@ -233,13 +248,16 @@ class Gates(SientiaMonitoring): self.debug(f'Filters: {filters}', metadata) comments = [] + + status = payload.status or {} + for fil, config in filters.items(): if fil not in mlflow_response_filter_functions: continue try: - if mlflow_response_filter_functions[fil](data, config): + if mlflow_response_filter_functions[fil](status, config): filter_output.append(config['policy']) - comments.append(data['content']['message']) + comments.append(status['message']) await self.send_notification_async( metadata=metadata, notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}', @@ -265,6 +283,9 @@ class Gates(SientiaMonitoring): 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') @@ -305,7 +326,10 @@ class Gates(SientiaMonitoring): self.info('Performing mlflow content gate...', metadata) filters = input_data['filters'] - data = DataFrame(input_data['data']) + + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) + gate_type = input_data['type'] path_priority = input_data['path_priority'] @@ -349,6 +373,9 @@ class Gates(SientiaMonitoring): ) self.info('Nothing was filtered by the mlflow content gate', metadata) + + del data + return None, 0, '' def get_prediction_store_policy( @@ -403,7 +430,7 @@ class Gates(SientiaMonitoring): return policy_type, int(policy_value) @activity.defn(name='format_transformed_data') - async def format_transformed_data(self, input_data: dict[str, Any]) -> dict: + async def format_transformed_data(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Format transformed data for storage and export operations. @@ -438,7 +465,8 @@ class Gates(SientiaMonitoring): self.info('Formatting transformed data...', metadata) - data = DataFrame(input_data['data']) + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) data['timestamp'] = data.index data = data.reset_index(drop=True) @@ -446,7 +474,13 @@ class Gates(SientiaMonitoring): data = data.melt(id_vars='timestamp', var_name='variable', value_name='value') data['model_id'] = model_id - return data.to_dict() + return await MinioDataFramePayload.from_dataframe( + dataframe=data, + minio_repo=self.minio_repository, + model_name=input_data['model_name'], + operation='transform', + workflow_metadata=metadata + ) @activity.defn(name='format_prediction') async def format_prediction(self, input_data: dict[str, Any]) -> dict: @@ -477,7 +511,8 @@ class Gates(SientiaMonitoring): prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) - data = DataFrame(input_data['data']) + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) # Create timestamp column from index and reset index data['timestamp'] = data.index @@ -566,6 +601,7 @@ class Gates(SientiaMonitoring): 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: """ @@ -636,45 +672,6 @@ class Gates(SientiaMonitoring): return report.to_dict() - @activity.defn(name='get_last_timestamp') - async def get_last_timestamp(self, input_data: dict[str, Any]) -> str: - """ - Extract the most recent timestamp from prediction data. - - This method analyzes prediction data to find the latest timestamp, - enabling incremental processing and data continuity tracking. - It handles empty datasets gracefully by returning the current time - as a fallback timestamp. - - The method is essential for: - 1. Incremental data processing workflows - 2. Data continuity validation - 3. Timestamp-based data loading optimization - 4. Workflow execution tracking - - Args: - input_data (dict): Input data containing: - - data (dict[str, Any]): Prediction data to analyze - - Returns: - str: Formatted timestamp string in UTC with timezone - """ - metadata = input_data['metadata'] - - self.info('Getting last timestamp...', metadata) - - data = DataFrame(input_data['data']) - - self.debug(f'Input data: {data.head(5).to_string()}', metadata) - - if data.empty: - return now().strftime(DATETIME_FORMAT_WITH_TZ) - - max_timestamp = max(data['timestamp'].values.tolist()) - - self.info(f'Last timestamp: {max_timestamp}', metadata) - - return max_timestamp @activity.defn(name='write_metrics') async def write_metrics(self, input_data: dict[str, Any]): diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 082e720..e593752 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,11 +1,15 @@ +from re import M from temporalio import activity, workflow +from laborious.utils.repository.minio_manager import MinioManager + with workflow.unsafe.imports_passed_through(): import traceback from typing import Any import numpy as np - from pandas import DataFrame, to_datetime + from io import BytesIO + from pandas import DataFrame, read_parquet, 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 @@ -19,11 +23,12 @@ with workflow.unsafe.imports_passed_through(): ) from sientia_do.utils.formatters import create_sample_dict - from laborious.utils.repository.minio_repository import MinioRepository + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + from sientia_do.repository.minio_repository import MinioRepository from laborious.utils.repository.model_repository import MLFlowRepository -class MLFlow(SientiaMonitoring): +class MLFlow(MinioManager): """ MLFlow integration activities for model inference operations. @@ -47,11 +52,11 @@ class MLFlow(SientiaMonitoring): mlflow_host: str, mlflow_port: int, mlflow_username: str, - minio_config: dict[str, Any], mlflow_password: str, - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, + 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. @@ -67,7 +72,7 @@ class MLFlow(SientiaMonitoring): Raises: Exception: If MLFlowRepository initialization fails """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) + 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 @@ -82,32 +87,17 @@ class MLFlow(SientiaMonitoring): metrics_controller, ) - if not hasattr(self, 'minio_repository'): - self.minio_repository: MinioRepository | None = None - - if self.minio_repository is None: - self.minio_repository = MinioRepository( - logger=logger, - notification_handler=notification_handler, - minio_endpoint_url=minio_config['endpoint_url'], - minio_access_key=minio_config['access_key'], - minio_secret_key=minio_config['secret_key'], - minio_region_name=minio_config['region_name'], - minio_default_bucket=minio_config['default_bucket'], - metrics_controller=metrics_controller, - ) - def close(self) -> None: """ Close the MLFlow activity and clean up resources. """ - SientiaMonitoring.shutdown(self) + MinioManager.close(self) def __del__(self): self.close() @activity.defn(name='request_transform') - async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Transform input data using MLFlow models. @@ -138,7 +128,10 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Transforming data...', metadata) - data = DataFrame(input_data['data']) + + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) + model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -176,10 +169,29 @@ class MLFlow(SientiaMonitoring): self.info('Data transformed successfully', metadata) - return response_data + 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, + ) + + 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, + }, + ) @activity.defn(name='request_predict') - async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]: + async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ Execute predictions using MLFlow models. @@ -210,7 +222,10 @@ class MLFlow(SientiaMonitoring): """ metadata = input_data['metadata'] self.info('Predicting data...', metadata) - data = DataFrame(input_data['data']) + + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) + model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) @@ -236,7 +251,28 @@ class MLFlow(SientiaMonitoring): self.info('Data predicted successfully', metadata) - return response_data + + 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, + ) + + 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, + } + ) + @activity.defn(name='retrain_model') async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]: @@ -275,14 +311,24 @@ class MLFlow(SientiaMonitoring): raise ValueError('Minio repository not initialized') metadata = input_data['metadata'] - object_key = input_data['object_key'] - - self.info(f'Loading retrain data from Key: {object_key}', metadata) try: - data = await self.minio_repository.get_parquet_as_dataframe( - object_key=object_key, metadata=metadata - ) + if 'data' in input_data: + # New path: payload-based retrain input (inline or MinIO offloaded). + data = await MinioDataFramePayload.dataframe_from_wire( + input_data['data'], + self.minio_repository, + metadata, + ) + else: + # Backward compatibility: legacy query_to_minio contract. + object_key = input_data['object_key'] + self.info(f'Loading retrain data from Key: {object_key}', metadata) + file_bytes = await self.minio_repository.download_file( + object_name=object_key, + metadata=metadata, + ) + data = read_parquet(BytesIO(file_bytes)) except Exception as e: trace = traceback.format_exc() await self.send_notification_async( diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 0b14397..d0f57d9 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -1,8 +1,15 @@ +import json 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 pickle import traceback + from datetime import timedelta + from io import BytesIO + from os import getenv from typing import Any import pandas as pd @@ -13,15 +20,20 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.temporal.activities.postgres import Postgres from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now - from laborious.utils.repository.minio_repository import MinioRepository + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + from sientia_do.repository.minio_repository import MinioRepository + +_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'}) -class Storage(Postgres): +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, @@ -31,12 +43,15 @@ class Storage(Postgres): dbname: str, min_connections: int, max_connections: int, - minio_config: dict[str, Any], - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, + retention_hours: int = 24, + minio_repository: MinioRepository | None = None, + logger: Logger | None = None, + notification_handler: NotificationHandler | None = None, + metrics_controller: MetricsController | None = None, ): - super().__init__( + self.retention_hours = retention_hours + Postgres.__init__( + self, host=host, port=port, user=user, @@ -49,20 +64,144 @@ class Storage(Postgres): metrics_controller=metrics_controller, ) - if not hasattr(self, 'minio_repository'): - self.minio_repository: MinioRepository | None = None + 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: - self.minio_repository = MinioRepository( - logger=logger, - notification_handler=notification_handler, - minio_endpoint_url=minio_config['endpoint_url'], - minio_access_key=minio_config['access_key'], - minio_secret_key=minio_config['secret_key'], - minio_region_name=minio_config['region_name'], - minio_default_bucket=minio_config['default_bucket'], - metrics_controller=metrics_controller, + 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', + ) + + @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 = 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', {}) + prefix = input_data['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: + await self.send_notification_async( + metadata=metadata, + notification_id='CLEANUP_MINIO_OBJECTS_EXPIRED', + message='MinIO objects cleaned up successfully', + block='cleanup_minio_objects_expired', + level=NotificationLevel.INFO, + attachment_content=json.dumps(report), + ) + + return report + + @activity.defn(name='query_to_minio') async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]: @@ -83,11 +222,18 @@ class Storage(Postgres): raise ValueError('Minio repository not initialized') metadata = input_data.get('metadata', {}) + model_name = input_data.get('model_name') or metadata.get('model_name') or 'unknown' object_prefix = input_data.get('object_prefix', 'datasets/retrain') timestamp = now().strftime(DATETIME_FORMAT_FILENAME) - object_name = f'{object_prefix}_{timestamp}.parquet' - uri = f's3://{self.minio_repository.minio_bucket}/{object_name}' + # Keep a stable model-level layout for minimal_retrain: + # training_datasets// + # Sanitize object_prefix to avoid extra subdirectories in the relative key. + safe_prefix = str(object_prefix).strip().strip('/').replace('/', '_') + filename = f'{safe_prefix}_{timestamp}.parquet' + relative_key = f'training_datasets/{model_name}/{filename}' + bucket = getattr(self.minio_repository, 'bucket', 'streamlit-connectors') + uri = f's3://{bucket}/{relative_key}' try: data = await self.load_custom_query(input_data) @@ -98,12 +244,20 @@ class Storage(Postgres): # Ensure we have a DataFrame data = pd.DataFrame(data) - # Write parquet to memory and upload via persistent client - await self.minio_repository.store_dataframe_as_parquet( - dataframe=data, uri=uri, object_name=object_name, metadata=metadata + # Convert DataFrame -> parquet bytes, then upload using the new MinIO interface. + parquet_buffer = BytesIO() + data.to_parquet(parquet_buffer, engine='pyarrow', index=True) + file_bytes = parquet_buffer.getvalue() + + upload_result = await self.minio_repository.upload_file( + file_bytes=file_bytes, + relative_key=relative_key, + metadata=metadata, ) - return {'success': True, 'object_key': object_name, 'uri': uri} + object_key_full = upload_result.get('minio_object_name', relative_key) + uri = f's3://{bucket}/{object_key_full}' + return {'success': True, 'object_key': object_key_full, 'uri': uri} except Exception as e: trace = traceback.format_exc() await self.send_notification_async( @@ -121,18 +275,8 @@ class Storage(Postgres): def close(self) -> None: """Close Storage resources (MinIO client and Postgres engine).""" - try: - if hasattr(self, 'minio_repository') and self.minio_repository is not None: - try: - self.minio_repository.close() - finally: - self.minio_repository = None - finally: - # Ensure Postgres resources are disposed as well - try: - super().close() - except Exception: - self.logger.error('Error closing Postgres resources') + Postgres.close(self) + MinioManager.close(self) def __del__(self): self.close() diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index c151377..5e3f186 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -88,4 +88,5 @@ def build_minio_config() -> dict[str, Any]: 'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'), 'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'), 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'), + 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), } 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..5ed7dac --- /dev/null +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -0,0 +1,230 @@ +""" +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 dataclasses import dataclass, field +from datetime import datetime +from io import BytesIO +from os import getenv +from typing import Any, Hashable, Literal + +from pandas import DataFrame, read_parquet + +from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now +from sientia_do.repository.minio_repository import MinioRepository + +# 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(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1.5')) * 1024 * 1024 + +# Relative prefix used for storing offloaded training datasets in MinIO. +# It is also the root directory for retention cleanup listing. +TRAINING_DATASETS_PREFIX = 'training_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 `training_datasets/` + # - Filename follows the retention-parsing pattern + basename = f'{model_name}-{operation}-{timestamp}.parquet' + model_dir = model_name.strip().strip('/') + prefix = f'{TRAINING_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 estimate_size_bytes(df: DataFrame) -> 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: + return len(pickle.dumps(df.to_dict())) + except Exception: + return len(pickle.dumps(df)) + + @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 + + @staticmethod + def is_offloaded_dict(payload: dict[str, Any]) -> bool: + """ + Return True if the dict represents a MinIO-backed payload without inline data. + + Args: + payload: Flat dict possibly produced by to_dict() / from_dataframe_to_dict(). + + Return: + bool: True when object_key is set and inline data is absent. + """ + if not payload.get('object_key'): + return False + return payload.get('data') is None + + @staticmethod + 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 + + @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, + ) -> '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 not dataframe or dataframe.empty: + return cls(data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status) + + last_timestamp = max(dataframe['timestamp'].values.tolist()) + + if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES: + return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp) + + timestamp = now().strftime(DATETIME_FORMAT_FILENAME) + object_key, object_prefix = _build_object_key(model_name, operation, timestamp) + + # 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 + + return cls( + data=None, + bucket=bucket, + object_key=object_key_full, + object_prefix=object_prefix, + uri=uri, + last_timestamp=last_timestamp, + ) + + async def retrieve(self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | 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: + return DataFrame(self.data) + + if self.data is None and self.object_key is None: + return DataFrame() + + file_bytes = await minio_repo.download_file( + object_name=self.object_key, metadata=workflow_metadata) + df = read_parquet(BytesIO(file_bytes)) + return df diff --git a/laborious/utils/repository/minio_manager.py b/laborious/utils/repository/minio_manager.py new file mode 100644 index 0000000..836862b --- /dev/null +++ b/laborious/utils/repository/minio_manager.py @@ -0,0 +1,26 @@ +from sientia_do.observability.logger import Logger +from sientia_do.observability.metrics_controller import MetricsController +from sientia_do.notifications.handlers import NotificationHandler +from sientia_do.repository.minio_repository import MinioRepository +from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + +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) \ No newline at end of file diff --git a/laborious/utils/repository/minio_repository.py b/laborious/utils/repository/minio_repository.py deleted file mode 100644 index fa9cf0f..0000000 --- a/laborious/utils/repository/minio_repository.py +++ /dev/null @@ -1,215 +0,0 @@ -""" -MinIO repository utilities. - -This module provides a lightweight repository around a MinIO/S3-compatible -object storage using boto3. It supports creating buckets on demand and -storing/loading pandas DataFrames in Parquet format. -""" - -import time -from io import BytesIO -from typing import Any - -import boto3 -from botocore.config import Config -from botocore.exceptions import ClientError -from pandas import DataFrame, read_parquet -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 laborious import metrics - - -class MinioRepository(SientiaMonitoring): - """ - Repository for interacting with a MinIO (S3-compatible) object storage. - - This class encapsulates a reusable `boto3` S3 client and convenience - helpers to persist and retrieve pandas DataFrames as Parquet files. - - Attributes: - storage_options (dict): Options compatible with pandas s3fs usage. - minio_bucket (str): Default bucket name used for operations. - minio_endpoint_url (str): MinIO endpoint URL. - minio_region_name (str): MinIO region name. - s3_client (Any): Reusable S3 client from `boto3`. - logger (Logger): Observability logger. - notification_handler (NotificationHandler): Notifications handler. - """ - - def __init__( - self, - minio_endpoint_url: str, - minio_access_key: str, - minio_secret_key: str, - minio_region_name: str, - minio_default_bucket: str, - logger: Logger, - notification_handler: NotificationHandler, - metrics_controller: MetricsController, - ): - """Initialize the repository and S3 client. - - Args: - minio_endpoint_url (str): MinIO endpoint URL. - minio_access_key (str): Access key (AK). - minio_secret_key (str): Secret key (SK). - minio_region_name (str): Region name for the client. - minio_default_bucket (str): Default bucket name to operate on. - logger (Logger): Logger instance for structured logs. - notification_handler (NotificationHandler): Notification handler. - """ - SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) - # MinIO settings shared with pandas s3fs - self.storage_options = { - 'key': minio_access_key, - 'secret': minio_secret_key, - 'client_kwargs': {'endpoint_url': minio_endpoint_url}, - } - self.minio_bucket = minio_default_bucket - self.minio_endpoint_url = minio_endpoint_url - self.minio_region_name = minio_region_name - - logger.info( - f'Connecting to MinIO at {self.minio_endpoint_url}, default bucket: {self.minio_bucket}' - ) - - # Reusable MinIO client - self.s3_client: Any = boto3.client( - 's3', - endpoint_url=self.minio_endpoint_url, - aws_access_key_id=self.storage_options['key'], - aws_secret_access_key=self.storage_options['secret'], - region_name=self.minio_region_name, - config=Config( - signature_version='s3v4', - s3={'addressing_style': 'path'}, - retries={'max_attempts': 5, 'mode': 'standard'}, - connect_timeout=5, - read_timeout=120, - ), - ) - - def close(self): - """Close the underlying S3 client.""" - self.s3_client.close() - - async def create_bucket(self, metadata: dict[str, Any]) -> None: - core_labels = { - **self.get_core_labels(metadata, operation_type='create_bucket'), - 'bucket_name': self.minio_bucket, - 'object_name': '-', - } - self.info(f"Creating bucket '{self.minio_bucket}'", metadata) - - start_time = time.time() - try: - self.s3_client.create_bucket(Bucket=self.minio_bucket) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels) - - async def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None: - """Ensure the default bucket exists; create it if missing. - - Args: - metadata (dict[str, Any]): Metadata used for structured logging. - """ - self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata) - core_labels = { - **self.get_core_labels(metadata, operation_type='head_bucket'), - 'bucket_name': self.minio_bucket, - 'object_name': '-', - } - self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata) - - start_time = time.time() - - try: - self.s3_client.head_bucket(Bucket=self.minio_bucket) - except ClientError: - await self.create_bucket(metadata) - - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - else: - await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels) - - async def store_dataframe_as_parquet( - self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any] - ): - """Persist a DataFrame as a Parquet object in the default bucket. - - Args: - dataframe (DataFrame): DataFrame to persist. - uri (str): Human-friendly URI used for logging context. - object_name (str): Object key (path/key within the bucket). - metadata (dict[str, Any]): Metadata used for structured logging. - """ - await self.ensure_bucket_exists(metadata) - - self.info(f'Storing dataframe as parquet in {uri}', metadata) - - buffer = BytesIO() - dataframe.to_parquet(buffer, engine='pyarrow', index=True) - buffer.seek(0) - - core_labels = { - **self.get_core_labels(metadata, operation_type='put_object'), - 'bucket_name': self.minio_bucket, - 'object_name': object_name, - } - start_time = time.time() - try: - self.s3_client.put_object( - Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue() - ) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels) - - self.info(f'Dataframe stored as parquet in {uri}', metadata) - - async def get_parquet_as_dataframe( - self, object_key: str, metadata: dict[str, Any] - ) -> DataFrame: - """Load a Parquet object from the default bucket into a DataFrame. - - Args: - object_key (str): Object key to retrieve from the bucket. - metadata (dict[str, Any]): Metadata used for structured logging. - - Returns: - DataFrame: Loaded DataFrame. - """ - self.info(f'Getting parquet as dataframe from {object_key}', metadata) - - core_labels = { - **self.get_core_labels(metadata, operation_type='get_object'), - 'bucket_name': self.minio_bucket, - 'object_name': object_key, - } - start_time = time.time() - try: - response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key) - except Exception as e: - await self.emit_metric(metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=core_labels) - raise e - - await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels) - - # Read the content into a BytesIO buffer to support seek operations - buffer = BytesIO(response['Body'].read()) - return read_parquet(buffer) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 66862e9..b5f9e9b 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1134,7 +1134,7 @@ class MLFlowRepository(SientiaMonitoring): 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. @@ -1196,7 +1196,7 @@ class MLFlowRepository(SientiaMonitoring): transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return {'success': True, 'content': transformed_data.to_dict()} + return {'success': True, 'content': transformed_data} except Exception as e: return { @@ -1290,7 +1290,7 @@ class MLFlowRepository(SientiaMonitoring): predict_data.index = input_index predict_data['response_time'] = (end_time - start_time).total_seconds() - return {'success': True, 'content': predict_data.to_dict()} + return {'success': True, 'content': predict_data} except Exception as e: return { diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 7281254..8c63b34 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -153,6 +153,7 @@ async def main(): other_workflows=[], activities=[ activities.load_custom_query, + activities.load_query_with_minio_offload, activities.query_to_minio, activities.retrain_model, activities.update_production_model, @@ -202,8 +203,9 @@ async def main(): activities.get_last_timestamp, # OPC activities.write_opc_data, - # Postgres - activities.load_custom_query, + # Postgres / MinIO offload + activities.load_query_with_minio_offload, + activities.cleanup_minio_objects_expired, activities.repeat_last_prediction, activities.export_data_to_postgres, activities.write_metrics, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 02878cf..72e11d7 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -73,26 +73,25 @@ class MinimalRetrain: model_config = input_data.get('model_config', {}) storage_result = await workflow.execute_activity_method( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': model_name, - 'object_prefix': f'retrain_datasets/{model_name}/data', }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=600), ) - if not storage_result['success']: + if isinstance(storage_result, dict) and storage_result.get('success') is False: return experiment_response = await workflow.execute_activity_method( Activities.retrain_model, { **metadata, - 'object_key': storage_result['object_key'], + 'data': storage_result, 'model_name': model_name, 'model_config': model_config, }, diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 42be048..427a470 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -83,13 +83,14 @@ class PredictionsBatch: } } - # Load data using custom query - data = await workflow.execute_local_activity_method( - Activities.load_custom_query, + # 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), diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 684f7a5..c9127c5 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -105,6 +105,7 @@ class FormatAndExportPrediction: '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), @@ -118,13 +119,14 @@ class FormatAndExportPrediction: **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_data_to_postgres, + Activities.export_payload_to_postgres, { **metadata, 'schema': input_data['schema'], diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 991d262..e307ab2 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -2,11 +2,13 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): from datetime import timedelta + from collections.abc import Callable 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='subworkflow.prediction_process') @@ -37,6 +39,8 @@ class PredictionProcess: 8. Export Delegation: Delegates to FormatAndExportPrediction workflow """ + cleanup_prefixes: set[str] = set() + @workflow.run async def run(self, input_data: dict[str, Any]): """ @@ -87,13 +91,39 @@ class PredictionProcess: model_config = input_data.get('model_config', {}) save_transform = input_data.get('save_transform', True) - # Get last timestamp for incremental processing - last_timestamp = await workflow.execute_local_activity_method( - Activities.get_last_timestamp, - {**metadata, 'data': data}, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=1), - ) + prefix = data.cleanup_prefix() + + try: + await self._run_prediction_pipeline( + input_data, + metadata, + data, + model_id, + model_name, + model_config, + save_transform, + ) + finally: + if self.cleanup_prefixes: + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'prefix': prefix}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) + + async def _run_prediction_pipeline( + self, + input_data: dict[str, Any], + metadata: dict[str, Any], + data: Any, + model_id: Any, + model_name: str, + model_config: dict[str, Any], + save_transform: bool + ) -> None: + + last_timestamp = data.last_timestamp # Apply input data quality gates gate_input = { @@ -119,7 +149,12 @@ class PredictionProcess: # Request MLFlow model transformation response_data = await workflow.execute_local_activity_method( Activities.request_transform, - {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, + { + **metadata, + 'data': data, + 'model_name': model_name, + 'model_config': model_config + }, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=5), ) @@ -221,7 +256,7 @@ class PredictionProcess: async def path_flag_handler( self, - data: dict, + data: MinioDataFramePayload, path_flag: str, input_data: dict, confidence: int, diff --git a/requirements.txt b/requirements.txt index c9f111c..4c30e62 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.8.2 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.2 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.7 prometheus-client botocore diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index dd7a0c5..46ae909 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -38,14 +38,13 @@ def test___init__(mock_minio_repository, mock_mlflow_repository): ) mock_minio_repository.assert_called_once_with( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', logger=ANY, notification_handler=ANY, - minio_endpoint_url='http://localhost:9000', - minio_access_key='minio', - minio_secret_key='minio123', - minio_region_name='us-east-1', - minio_default_bucket='test', metrics_controller=ANY, + bucket='test', ) @@ -96,10 +95,15 @@ metadata = { @mark.asyncio -@patch('laborious.activities.mlflow.DataFrame') +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', + new_callable=AsyncMock, +) @patch('laborious.activities.mlflow.max') -async def test_request_transform_success(mock_max, mock_dataframe, mlflow): +async def test_request_transform_success(mock_max, mock_dataframe_from_wire, mlflow): mock_max.return_value = '2024-01-02' + data_mock = MagicMock() + mock_dataframe_from_wire.return_value = data_mock # Mock input data input_data = { **metadata, @@ -149,37 +153,41 @@ async def test_request_transform_success(mock_max, mock_dataframe, mlflow): expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']} mlflow.model_monitoring_repository.transform.return_value = expected_response - mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value - mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value + data_mock.sort_values.return_value = data_mock + data_mock.drop_duplicates.return_value = data_mock + data_mock.pivot.return_value = data_mock # Call the method response_data = await mlflow.request_transform(input_data) # Verify the data was correctly transformed - mock_dataframe.assert_called_once_with(input_data['data']) - mock_dataframe.return_value.pivot.assert_called_once_with( + data_mock.pivot.assert_called_once_with( index='timestamp', columns='variable', values='value' ) - mock_dataframe = mock_dataframe.return_value.pivot.return_value - mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True) + data_mock.fillna.assert_called_once_with(np.nan, inplace=True) # mock_dataframe.reset_index.assert_called_once() - mock_dataframe.columns.name = None + data_mock.columns.name = None # Verify the response assert response_data == expected_response # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.transform.assert_called_once_with( - 'test_model', mock_dataframe, {}, metadata['metadata'] + 'test_model', data_mock, {}, metadata['metadata'] ) @mark.asyncio -@patch('laborious.activities.mlflow.DataFrame') +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', + new_callable=AsyncMock, +) @patch('laborious.activities.mlflow.to_datetime') @patch('laborious.activities.mlflow.max') -async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow): +async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe_from_wire, mlflow): mock_max.return_value = '2024-01-02' + data_mock = MagicMock() + mock_dataframe_from_wire.return_value = data_mock # Mock input data input_data = { **metadata, @@ -203,22 +211,21 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo # Call the method response_data = await mlflow.request_predict(input_data) - mock_dataframe.assert_called_once_with(input_data['data']) - mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True) - mock_dataframe.return_value.__setitem__.assert_any_call( + data_mock.replace.assert_called_once_with(np.nan, None, inplace=True) + data_mock.__setitem__.assert_any_call( 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value ) - mock_dataframe.return_value.__setitem__.assert_any_call( + data_mock.__setitem__.assert_any_call( 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value ) mock_to_datetime.assert_called_once_with( - mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ + data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) mock_to_datetime.assert_called_once_with( - mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ + data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ ) mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) @@ -227,20 +234,22 @@ async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflo # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.predict.assert_called_once_with( - 'test_model', mock_dataframe.return_value, {}, metadata['metadata'] + 'test_model', data_mock, {}, metadata['metadata'] ) @mark.asyncio +@patch('laborious.activities.mlflow.read_parquet') @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow): +async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mock_read_parquet, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', } - mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock() + mlflow.minio_repository.download_file.return_value = b'parquet-bytes' + mock_read_parquet.return_value = MagicMock() response = await mlflow.retrain_model( { @@ -255,7 +264,7 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl } ) - raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value + raw_data = mock_read_parquet.return_value timestamp = raw_data.__getitem__.return_value.max.return_value @@ -308,15 +317,47 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl @mark.asyncio +@patch('laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', new_callable=AsyncMock) @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow): +async def test_retrain_model_success_with_payload_data( + mock_to_datetime, mock_dataframe_from_wire, mlflow +): + mlflow.model_monitoring_repository.retrain_model.return_value = { + 'success': True, + 'experiment': 'test_experiment', + 'message': 'Model retrained successfully.', + } + mock_dataframe_from_wire.return_value = MagicMock() + + response = await mlflow.retrain_model( + { + **metadata, + 'data': {'data': {'a': [1]}}, + 'model_name': 'test_model', + 'model_config': { + 'target': 'target', + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + }, + } + ) + + assert response['success'] is True + mlflow.minio_repository.download_file.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.mlflow.read_parquet') +@patch('laborious.activities.mlflow.to_datetime') +async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mock_read_parquet, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': False, 'traceback': 'test_traceback', 'message': 'Model retrained failed.', } - mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock( + mlflow.minio_repository.download_file.return_value = b'parquet-bytes' + mock_read_parquet.return_value = MagicMock( columns=['variable', 'timestamp', 'value', 'created_at'] ) @@ -333,7 +374,7 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow) } ) - raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value + raw_data = mock_read_parquet.return_value timestamp = raw_data.__getitem__.return_value.max.return_value @@ -398,7 +439,7 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow) @mark.asyncio async def test_retrain_model_data_error(mlflow): - mlflow.minio_repository.get_parquet_as_dataframe.side_effect = Exception( + mlflow.minio_repository.download_file.side_effect = Exception( 'Error loading retrain data' ) diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index 411e22e..963dbf2 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -1,4 +1,5 @@ import datetime +import os from unittest.mock import ANY, AsyncMock, MagicMock, patch from pytest import fixture, mark, raises @@ -68,14 +69,13 @@ def test___init___not_hasattr(mock_minio_repository): assert isinstance(storage, Postgres) mock_minio_repository.assert_called_once_with( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', logger=logger, notification_handler=notification_handler, - minio_endpoint_url='localhost:9000', - minio_access_key='minio', - minio_secret_key='minio123', - minio_region_name='us-east-1', - minio_default_bucket='test', metrics_controller=metrics_controller, + bucket='test', ) @@ -106,14 +106,13 @@ def test___init___none_minio_repository(mock_minio_repository, storage): ) mock_minio_repository.assert_called_once_with( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', logger=logger, notification_handler=notification_handler, - minio_endpoint_url='localhost:9000', - minio_access_key='minio', - minio_secret_key='minio123', - minio_region_name='us-east-1', - minio_default_bucket='test', metrics_controller=metrics_controller, + bucket='test', ) @@ -169,30 +168,38 @@ async def test_query_to_minio_success(now, dataframe, storage): data = [{'a': 1}, {'a': 2}, {'a': 3}] storage.load_custom_query = AsyncMock(return_value=data) now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0) - storage.minio_repository.store_dataframe_as_parquet = AsyncMock() - storage.minio_repository.minio_bucket = 'test' + storage.minio_repository.upload_file = AsyncMock( + return_value={ + 'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' + } + ) + storage.minio_repository.bucket = 'test' result = await storage.query_to_minio({'object_prefix': 'test', **metadata}) dataframe.assert_called_once_with(data) - storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with( - dataframe=dataframe.return_value, - uri='s3://test/test_2024-01-01_00-00-00.parquet', - object_name='test_2024-01-01_00-00-00.parquet', + storage.minio_repository.upload_file.assert_called_once_with( + file_bytes=ANY, + relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet', metadata=metadata['metadata'], ) assert result['success'] is True - assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet' - assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet' + assert ( + result['object_key'] + == 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' + ) + assert ( + result['uri'] + == 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' + ) @mark.asyncio async def test_query_to_minio_error(storage): storage.send_notification = MagicMock() storage.send_notification_async = AsyncMock() - storage.minio_repository.store_dataframe_as_parquet = AsyncMock() storage.load_custom_query = AsyncMock(side_effect=Exception('test')) result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'}) @@ -222,3 +229,80 @@ def test___del__(storage): storage.__del__() storage.close.assert_called_once() + + +def test_estimate_payload_size_bytes(storage): + assert storage._estimate_payload_size_bytes({'x': 1}) > 0 + + +@mark.asyncio +async def test_load_query_with_minio_offload_no_rows(storage): + storage.load_custom_query = AsyncMock(return_value=None) + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} + ) + assert result['success'] is False + + +@mark.asyncio +async def test_load_query_with_minio_offload_inline(storage): + storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'my-model', 'key_prefix': 'predictions/s'} + ) + assert result.get('success') is True + assert 'data' in result + assert result.get('object_key') is None + + +@mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.MinioDataFramePayload.estimate_size_bytes') +async def test_load_query_with_minio_offload_minio(mock_estimate, storage): + mock_estimate.return_value = 10**9 + storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) + storage.minio_repository.upload_file = AsyncMock( + return_value={ + 'minio_object_name': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet' + } + ) + storage.minio_repository.bucket = 'test' + + fixed = datetime.datetime(2024, 1, 15, 12, 30, 45) + with patch('laborious.utils.models.minio_dataframe_payload.now', return_value=fixed): + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} + ) + + assert result.get('success') is True + assert result.get('data') is None + assert ( + result['object_key'] + == 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet' + ) + storage.minio_repository.upload_file.assert_called_once() + + +@mark.asyncio +@patch.dict(os.environ, {'SIENTIA_MINIO_RETENTION_HOURS': '1'}) +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock( + return_value=[ + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet', + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2025-01-10_12-00-00.parquet', + ] + ) + storage.minio_repository.delete_file = AsyncMock() + storage.send_notification_async = AsyncMock() + + result = await storage.cleanup_minio_objects_expired( + {**metadata, 'prefixes': ['training_datasets/m']} + ) + + assert result['success'] is True + assert result['deleted_count'] == 1 + storage.minio_repository.delete_file.assert_called_once_with( + object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet', + metadata=metadata['metadata'], + ) diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py new file mode 100644 index 0000000..71aeda0 --- /dev/null +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -0,0 +1,54 @@ +from datetime import datetime + +from pytest import mark + +from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + + +def test_parse_object_timestamp_hyphenated_model(): + key = 'predictions/sched/my-long-model-initial-2024-06-15_10-30-45.parquet' + ts = MinioDataFramePayload.parse_object_timestamp(key) + assert ts == datetime(2024, 6, 15, 10, 30, 45) + + +def test_parse_object_timestamp_transform(): + key = 'p/m-transform-2024-01-02_03-04-05.parquet' + ts = MinioDataFramePayload.parse_object_timestamp(key) + assert ts == datetime(2024, 1, 2, 3, 4, 5) + + +def test_parse_object_timestamp_invalid(): + assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None + + +def test_is_offloaded_dict_true_false(): + assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': None}) is True + assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': {}}) is False + assert MinioDataFramePayload.is_offloaded_dict({'data': {}}) is False + + +def test_cleanup_prefix_from_payload_dict(): + p = { + 'object_key': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-01_00-00-00.parquet', + 'bucket': 'b', + 'data': None, + } + assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'training_datasets/m' + + +def test_cleanup_prefix_from_explicit_object_prefix(): + p = {'object_key': 'x.parquet', 'object_prefix': 'my/prefix', 'data': None} + assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'my/prefix' + + +@mark.asyncio +async def test_resolve_dict_if_offloaded_noop(): + d = {'success': True, 'data': {'a': [1]}} + out = await MinioDataFramePayload.resolve_dict_if_offloaded(d, None, {}) + assert out is d + + +@mark.asyncio +async def test_dataframe_from_wire_list(): + df = await MinioDataFramePayload.dataframe_from_wire([{'a': 1}], None, {}) + assert list(df.columns) == ['a'] diff --git a/tests/laborious/utils/repository/test_minio_repository.py b/tests/laborious/utils/repository/test_minio_repository.py deleted file mode 100644 index 89d287e..0000000 --- a/tests/laborious/utils/repository/test_minio_repository.py +++ /dev/null @@ -1,245 +0,0 @@ -from unittest.mock import ANY, AsyncMock, MagicMock, patch - -from botocore.utils import ClientError -from pytest import fixture, mark, raises - -from laborious import metrics -from laborious.utils.repository.minio_repository import MinioRepository - - -@patch('laborious.utils.repository.minio_repository.boto3') -@patch('laborious.utils.repository.minio_repository.Config') -def test___init___(mock_config, mock_boto3): - minio_repository = MinioRepository( - minio_endpoint_url='localhost:9000', - minio_access_key='minio', - minio_secret_key='minio123', - minio_region_name='us-east-1', - minio_default_bucket='test', - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), - ) - - assert minio_repository.storage_options == { - 'key': 'minio', - 'secret': 'minio123', - 'client_kwargs': {'endpoint_url': 'localhost:9000'}, - } - assert minio_repository.minio_bucket == 'test' - assert minio_repository.minio_endpoint_url == 'localhost:9000' - assert minio_repository.minio_region_name == 'us-east-1' - - mock_config.assert_called_once_with( - signature_version='s3v4', - s3={'addressing_style': 'path'}, - retries={'max_attempts': 5, 'mode': 'standard'}, - connect_timeout=5, - read_timeout=120, - ) - - mock_boto3.client.assert_called_once_with( - 's3', - endpoint_url='localhost:9000', - aws_access_key_id='minio', - aws_secret_access_key='minio123', - region_name='us-east-1', - config=mock_config.return_value, - ) - - -@fixture -@patch('laborious.utils.repository.minio_repository.Config') -@patch('laborious.utils.repository.minio_repository.boto3') -def minio_repository(mock_boto3, mock_config): - minio_repository = MinioRepository( - minio_endpoint_url='localhost:9000', - minio_access_key='minio', - minio_secret_key='minio123', - minio_region_name='us-east-1', - minio_default_bucket='test', - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), - ) - - minio_repository.emit_metric = AsyncMock() - minio_repository.observe_lag = AsyncMock() - minio_repository.send_notification = MagicMock() - minio_repository.send_notification_async = AsyncMock() - - return minio_repository - - -def test_close(minio_repository): - minio_repository.close() - minio_repository.s3_client.close.assert_called_once() - - -@mark.asyncio -async def test_create_bucket_success(minio_repository): - await minio_repository.create_bucket({}) - - minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY - ) - - -@mark.asyncio -async def test_create_bucket_error(minio_repository): - minio_repository.s3_client.create_bucket.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.create_bucket({}) - - minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test') - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() - - -@mark.asyncio -async def test_ensure_bucket_exists_bucket_exists(minio_repository): - assert await minio_repository.ensure_bucket_exists({}) is None - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_COUNT, tags=ANY - ) - - -@mark.asyncio -async def test_ensure_bucket_exists_bucket_not_exists_create_success(minio_repository): - minio_repository.s3_client.head_bucket.side_effect = ClientError( - error_response={'Error': {'Code': '404'}}, operation_name='head_bucket' - ) - minio_repository.create_bucket = AsyncMock() - - assert await minio_repository.ensure_bucket_exists({}) is None - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - minio_repository.create_bucket.assert_called_once_with({}) - - minio_repository.observe_lag.assert_not_called() - minio_repository.emit_metric.assert_not_called() - - -@mark.asyncio -async def test_ensure_bucket_exists_bucket_not_exists_create_error(minio_repository): - minio_repository.s3_client.head_bucket.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.ensure_bucket_exists({}) - - minio_repository.s3_client.head_bucket.assert_called_once_with(Bucket='test') - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -async def test_store_dataframe_as_parquet_success(mock_bytesio, minio_repository): - input_data = MagicMock() - - minio_repository.ensure_bucket_exists = AsyncMock() - - await minio_repository.store_dataframe_as_parquet( - dataframe=input_data, uri='s3://test/test.parquet', object_name='test.parquet', metadata={} - ) - - minio_repository.ensure_bucket_exists.assert_called_once_with({}) - mock_bytesio.assert_called_once() - - input_data.to_parquet.assert_called_once_with( - mock_bytesio.return_value, engine='pyarrow', index=True - ) - mock_bytesio.return_value.seek.assert_called_once_with(0) - minio_repository.s3_client.put_object.assert_called_once_with( - Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value - ) - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_WRITE_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -async def test_store_dataframe_as_parquet_error(mock_bytesio, minio_repository): - input_data = MagicMock() - - minio_repository.ensure_bucket_exists = AsyncMock() - minio_repository.s3_client.put_object.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.store_dataframe_as_parquet( - dataframe=input_data, - uri='s3://test/test.parquet', - object_name='test.parquet', - metadata={}, - ) - - minio_repository.ensure_bucket_exists.assert_called_once_with({}) - mock_bytesio.assert_called_once() - - input_data.to_parquet.assert_called_once_with( - mock_bytesio.return_value, engine='pyarrow', index=True - ) - mock_bytesio.return_value.seek.assert_called_once_with(0) - minio_repository.s3_client.put_object.assert_called_once_with( - Bucket='test', Key='test.parquet', Body=mock_bytesio.return_value.getvalue.return_value - ) - - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -@patch('laborious.utils.repository.minio_repository.read_parquet') -async def test_get_parquet_as_dataframe_success(mock_read_parquet, mock_bytesio, minio_repository): - input_data = {'Body': MagicMock(read=MagicMock(return_value=b'test'))} - - minio_repository.s3_client.get_object.return_value = input_data - - output = await minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={}) - - minio_repository.s3_client.get_object.assert_called_once_with(Bucket='test', Key='test.parquet') - - mock_bytesio.assert_called_once_with(input_data['Body'].read.return_value) - mock_read_parquet.assert_called_once_with(mock_bytesio.return_value) - - assert output == mock_read_parquet.return_value - - minio_repository.observe_lag.assert_called_once_with(ANY, metrics.MINIO_READ_LAG, ANY) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_COUNT, tags=ANY - ) - - -@mark.asyncio -@patch('laborious.utils.repository.minio_repository.BytesIO') -@patch('laborious.utils.repository.minio_repository.read_parquet') -async def test_get_parquet_as_dataframe_error(mock_read_parquet, mock_bytesio, minio_repository): - minio_repository.s3_client.get_object.side_effect = ValueError('test') - - with raises(ValueError): - await minio_repository.get_parquet_as_dataframe(object_key='test.parquet', metadata={}) - - minio_repository.s3_client.get_object.assert_called_once_with( - Bucket='test', Key='test.parquet' - ) - minio_repository.emit_metric.assert_called_once_with( - metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=ANY - ) - minio_repository.observe_lag.assert_not_called() diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 22df274..4132537 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -17,6 +17,7 @@ metadata = { 'model_name': 'test_model', 'workflow_name': 'test_workflow', 'schema_name': 'test_schedule', + 'schedule_name': 'test_schedule', }, } @@ -101,6 +102,7 @@ async def test_run(workflow_mock, prediction_process): 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], + 'key_prefix': 'predictions/test_schedule', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -323,6 +325,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], + 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, @@ -423,6 +426,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], + 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, @@ -540,6 +544,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], + 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index bcbad2d..6413019 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -41,7 +41,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': True, 'object_key': 'test_object_key'}, + {'data': {'a': [1]}, 'success': True}, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -58,13 +58,12 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': input_data['model_name'], - 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -78,7 +77,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): Activities.retrain_model, { **metadata, - 'object_key': 'test_object_key', + 'data': {'data': {'a': [1]}, 'success': True}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, @@ -163,7 +162,7 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': False, 'object_key': 'test_object_key'}, + {'success': False, 'message': 'No data returned from query'}, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -178,13 +177,12 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim await minimal_retrain.run(input_data) workflow_mock.execute_activity_method.assert_called_once_with( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': input_data['model_name'], - 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -213,7 +211,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': True, 'object_key': 'test_object_key'}, + {'data': {'a': [1]}, 'success': True}, {'success': False, 'experiment': 'test_experiment'}, { 'success': True, @@ -230,13 +228,12 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.query_to_minio, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': input_data['model_name'], - 'object_prefix': f'retrain_datasets/{input_data["model_name"]}/data', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -250,7 +247,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim Activities.retrain_model, { **metadata, - 'object_key': 'test_object_key', + 'data': {'data': {'a': [1]}, 'success': True}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index c19d6ef..5505e89 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -24,7 +24,10 @@ metadata = { @mark.asyncio @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): - workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'} + workflow_mock.execute_activity_method.return_value = { + 'success': True, + 'data': {'col': ['test_data']}, + } input_data = { 'schedule_name': 'test_schedule', 'model_name': 'test_model', @@ -42,14 +45,16 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch await predictions_batch.run(input_data) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( - Activities.load_custom_query, + Activities.load_query_with_minio_offload, { **metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), + 'model_name': input_data['model_name'], + 'key_prefix': f"predictions/{input_data['schedule_name']}", }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -58,7 +63,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch ) prediction_input = { 'metadata': metadata, - 'data': {'data': 'test_data'}, + 'data': {'success': True, 'data': {'col': ['test_data']}}, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'transform_table_name': input_data['transform_table_name'], From 5d0d049082f631c78f6d22377c6b918c67c9c133 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 09:14:16 -0300 Subject: [PATCH 02/51] SIENTIAPDE-1712 Remove code validation script and refactor imports in activities and workflows - Deleted the `validate.sh` script, which was responsible for running code quality checks. - Cleaned up import statements in `activities.py`, `gates.py`, `mlflow.py`, and `storage.py` by removing unused imports and organizing them. - Refactored initialization methods in `MinioManager` and `MLFlow` classes for improved readability. - Updated various workflows to ensure compatibility with the new structure and removed unnecessary comments. - Enhanced test cases to accommodate changes in the activities and workflows, ensuring proper mocking of dependencies. --- laborious/activities/activities.py | 1 - laborious/activities/gates.py | 18 +- laborious/activities/mlflow.py | 37 +-- laborious/activities/storage.py | 19 +- .../utils/models/minio_dataframe_payload.py | 47 ++- laborious/utils/repository/minio_manager.py | 16 +- laborious/workflows/drift.py | 2 +- laborious/workflows/minimal_retrain.py | 4 +- laborious/workflows/simple_metrics.py | 2 +- .../sub_workflows/prediction_process.py | 11 +- tests/conftest.py | 49 ++- tests/laborious/activities/test_activities.py | 22 +- tests/laborious/activities/test_api.py | 54 ++++ tests/laborious/activities/test_gates.py | 296 ++++++++++------- tests/laborious/activities/test_mlflow.py | 303 +++++++++--------- .../activities/test_model_metrics.py | 119 +++++++ tests/laborious/activities/test_storage.py | 237 +++++++++----- .../models/test_minio_dataframe_payload.py | 212 ++++++++++-- .../utils/repository/test_model_repository.py | 149 ++++++++- .../laborious/utils/test_connectors_config.py | 2 + .../test_format_and_export_prediction.py | 24 +- .../subworkflows/test_prediction_process.py | 154 +++++---- .../workflows/test_minimal_retrain.py | 26 +- .../workflows/test_predictions_batch.py | 1 - validate.sh | 124 ------- 25 files changed, 1224 insertions(+), 705 deletions(-) delete mode 100755 validate.sh diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index f714a4f..39ee7d7 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -14,7 +14,6 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.model_metrics import ModelMetrics from laborious.activities.opc import OPC from laborious.activities.storage import Storage - class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index f226cc8..bd40b53 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -13,17 +13,15 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.models import NotificationLevel from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController - from sientia_do.observability.sientia_monitoring import SientiaMonitoring - from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now from sientia_do.utils.formatters import create_sample_dict from laborious import metrics - from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload 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] @@ -106,7 +104,9 @@ class Gates(MinioManager): Raises: Exception: If BaseActivity initialization fails """ - MinioManager.__init__(self, minio_repository, logger, notification_handler, metrics_controller) + MinioManager.__init__( + self, minio_repository, logger, notification_handler, metrics_controller + ) def close(self) -> None: """ @@ -238,7 +238,7 @@ class Gates(MinioManager): payload: MinioDataFramePayload = input_data['data'] data = await payload.retrieve(self.minio_repository, metadata) - + gate_type = input_data['type'] path_priority = input_data['path_priority'] @@ -326,7 +326,7 @@ class Gates(MinioManager): self.info('Performing mlflow content gate...', metadata) filters = input_data['filters'] - + payload: MinioDataFramePayload = input_data['data'] data = await payload.retrieve(self.minio_repository, metadata) @@ -375,7 +375,7 @@ class Gates(MinioManager): self.info('Nothing was filtered by the mlflow content gate', metadata) del data - + return None, 0, '' def get_prediction_store_policy( @@ -479,7 +479,7 @@ class Gates(MinioManager): minio_repo=self.minio_repository, model_name=input_data['model_name'], operation='transform', - workflow_metadata=metadata + workflow_metadata=metadata, ) @activity.defn(name='format_prediction') @@ -601,7 +601,6 @@ class Gates(MinioManager): 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: """ @@ -672,7 +671,6 @@ class Gates(MinioManager): return report.to_dict() - @activity.defn(name='write_metrics') async def write_metrics(self, input_data: dict[str, Any]): """ diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index e593752..82c664f 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,4 +1,3 @@ -from re import M from temporalio import activity, workflow from laborious.utils.repository.minio_manager import MinioManager @@ -8,13 +7,12 @@ with workflow.unsafe.imports_passed_through(): from typing import Any import numpy as np - from io import BytesIO - from pandas import DataFrame, read_parquet, to_datetime + 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.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository from sientia_do.temporal.constants import ( DATETIME_FORMAT, DATETIME_FORMAT_MS_WITH_TZ, @@ -24,7 +22,6 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.utils.formatters import create_sample_dict from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload - from sientia_do.repository.minio_repository import MinioRepository from laborious.utils.repository.model_repository import MLFlowRepository @@ -72,7 +69,9 @@ class MLFlow(MinioManager): Raises: Exception: If MLFlowRepository initialization fails """ - MinioManager.__init__(self, minio_repository, logger, notification_handler, metrics_controller) + 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 @@ -128,7 +127,7 @@ class MLFlow(MinioManager): """ metadata = input_data['metadata'] self.info('Transforming data...', metadata) - + payload: MinioDataFramePayload = input_data['data'] data = await payload.retrieve(self.minio_repository, metadata) @@ -251,7 +250,6 @@ class MLFlow(MinioManager): self.info('Data predicted successfully', metadata) - if not response_data.get('success', False): return await MinioDataFramePayload.from_dataframe( dataframe=None, @@ -270,10 +268,9 @@ class MLFlow(MinioManager): workflow_metadata=metadata, status={ 'success': True, - } + }, ) - @activity.defn(name='retrain_model') async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]: """ @@ -313,22 +310,10 @@ class MLFlow(MinioManager): metadata = input_data['metadata'] try: - if 'data' in input_data: - # New path: payload-based retrain input (inline or MinIO offloaded). - data = await MinioDataFramePayload.dataframe_from_wire( - input_data['data'], - self.minio_repository, - metadata, - ) - else: - # Backward compatibility: legacy query_to_minio contract. - object_key = input_data['object_key'] - self.info(f'Loading retrain data from Key: {object_key}', metadata) - file_bytes = await self.minio_repository.download_file( - object_name=object_key, - metadata=metadata, - ) - data = read_parquet(BytesIO(file_bytes)) + # Payload-based retrain input (inline dict or MinIO offloaded). + payload: MinioDataFramePayload = input_data['data'] + data = await payload.retrieve(self.minio_repository, metadata) + except Exception as e: trace = traceback.format_exc() await self.send_notification_async( diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index d0f57d9..9fda091 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -1,15 +1,14 @@ import json + 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 pickle import traceback from datetime import timedelta from io import BytesIO - from os import getenv from typing import Any import pandas as pd @@ -17,11 +16,11 @@ with workflow.unsafe.imports_passed_through(): 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 DATETIME_FORMAT_FILENAME, now from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload - from sientia_do.repository.minio_repository import MinioRepository _LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'}) @@ -64,10 +63,14 @@ class Storage(Postgres, MinioManager): metrics_controller=metrics_controller, ) - MinioManager.__init__(self, minio_repository, logger, notification_handler, 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: + 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. @@ -92,7 +95,9 @@ class Storage(Postgres, MinioManager): input_data, ) if not rows: - self.error('load_query_with_minio_offload failed: No data returned from query', metadata) + self.error( + 'load_query_with_minio_offload failed: No data returned from query', metadata + ) dataframe = None else: dataframe = pd.DataFrame(rows) @@ -200,8 +205,6 @@ class Storage(Postgres, MinioManager): ) return report - - @activity.defn(name='query_to_minio') async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]: diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 5ed7dac..4a51bb3 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -12,16 +12,16 @@ Otherwise, it is inlined as a Temporal-friendly ``dict``. import pickle import re -from dataclasses import dataclass, field +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, Hashable, Literal +from typing import Any, Literal from pandas import DataFrame, read_parquet - -from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now 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'}) @@ -30,7 +30,9 @@ _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(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1.5')) * 1024 * 1024 +OFFLOAD_THRESHOLD_BYTES = int( + float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1.5')) * 1024 * 1024 +) # Relative prefix used for storing offloaded training datasets in MinIO. # It is also the root directory for retention cleanup listing. @@ -79,7 +81,6 @@ class MinioDataFramePayload: object_prefix: str | None = None uri: str | None = None - @staticmethod def estimate_size_bytes(df: DataFrame) -> int: """ @@ -117,21 +118,6 @@ class MinioDataFramePayload: except ValueError: return None - @staticmethod - def is_offloaded_dict(payload: dict[str, Any]) -> bool: - """ - Return True if the dict represents a MinIO-backed payload without inline data. - - Args: - payload: Flat dict possibly produced by to_dict() / from_dataframe_to_dict(). - - Return: - bool: True when object_key is set and inline data is absent. - """ - if not payload.get('object_key'): - return False - return payload.get('data') is None - @staticmethod def cleanup_prefix(self) -> str | None: """ @@ -141,6 +127,12 @@ class MinioDataFramePayload: 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 not self.data != {}) or self.object_key is not None + @classmethod async def from_dataframe( cls, @@ -172,7 +164,9 @@ class MinioDataFramePayload: """ if not dataframe or dataframe.empty: - return cls(data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status) + return cls( + data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status + ) last_timestamp = max(dataframe['timestamp'].values.tolist()) @@ -207,7 +201,9 @@ class MinioDataFramePayload: last_timestamp=last_timestamp, ) - async def retrieve(self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | None = None) -> DataFrame: + async def retrieve( + self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | None = None + ) -> DataFrame: """ Load parquet from MinIO when object_key is set and populate inline data. @@ -221,10 +217,11 @@ class MinioDataFramePayload: if self.data is not None: return DataFrame(self.data) - if self.data is None and self.object_key is None: + if not self.has_data(): return DataFrame() file_bytes = await minio_repo.download_file( - object_name=self.object_key, metadata=workflow_metadata) + object_name=self.object_key, metadata=workflow_metadata + ) df = read_parquet(BytesIO(file_bytes)) return df diff --git a/laborious/utils/repository/minio_manager.py b/laborious/utils/repository/minio_manager.py index 836862b..17c85e6 100644 --- a/laborious/utils/repository/minio_manager.py +++ b/laborious/utils/repository/minio_manager.py @@ -1,14 +1,20 @@ +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.notifications.handlers import NotificationHandler -from sientia_do.repository.minio_repository import MinioRepository 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): + 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) @@ -23,4 +29,4 @@ class MinioManager(SientiaMonitoring): finally: self.minio_repository = None - SientiaMonitoring.shutdown(self) \ No newline at end of file + SientiaMonitoring.shutdown(self) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 769c613..b98fc56 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -44,7 +44,7 @@ class Drift: 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_local_activity_method( Activities.load_custom_query, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 72e11d7..f0cf0e8 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -84,8 +84,8 @@ class MinimalRetrain: start_to_close_timeout=timedelta(seconds=600), ) - if isinstance(storage_result, dict) and storage_result.get('success') is False: - return + if not storage_result.has_data(): + raise ValueError('No data returned from query') experiment_response = await workflow.execute_activity_method( Activities.retrain_model, diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index 438b98d..20b60ab 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -45,7 +45,7 @@ class SimpleMetrics: 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_local_activity_method( Activities.load_custom_query, diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index e307ab2..07e42e7 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -2,7 +2,6 @@ from temporalio import workflow with workflow.unsafe.imports_passed_through(): from datetime import timedelta - from collections.abc import Callable from typing import Any from sientia_do.temporal.policies import retry_policy @@ -120,9 +119,8 @@ class PredictionProcess: model_id: Any, model_name: str, model_config: dict[str, Any], - save_transform: bool + save_transform: bool, ) -> None: - last_timestamp = data.last_timestamp # Apply input data quality gates @@ -149,12 +147,7 @@ class PredictionProcess: # Request MLFlow model transformation response_data = await workflow.execute_local_activity_method( Activities.request_transform, - { - **metadata, - 'data': data, - 'model_name': model_name, - 'model_config': model_config - }, + {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, retry_policy=retry_policy, start_to_close_timeout=timedelta(minutes=5), ) diff --git a/tests/conftest.py b/tests/conftest.py index c0fad01..0cf3b48 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,3 +1,49 @@ +import os +import sys +from unittest.mock import MagicMock + +# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES to int at import-time. +# Tests must set it to a valid integer string to avoid import errors. +os.environ.setdefault('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1') + + +class DummyMinioDataFramePayload: + """ + Minimal payload double used by unit tests. + + The production workflow/gates expect a MinioDataFramePayload-like object with: + - async retrieve(minio_repo, workflow_metadata) -> DataFrame | dict + - has_data() -> bool + - cleanup_prefix() -> str | None + - last_timestamp: attribute + - status: attribute + """ + + def __init__( + self, + *, + retrieve_return=None, + has_data: bool = True, + cleanup_prefix: str | None = None, + last_timestamp: str = '2024-01-01', + status: dict | None = None, + ): + self._retrieve_return = retrieve_return + self._has_data = has_data + self._cleanup_prefix = cleanup_prefix + self.last_timestamp = last_timestamp + self.status = status + + async def retrieve(self, _minio_repo, _workflow_metadata=None): + return self._retrieve_return + + def has_data(self) -> bool: + return self._has_data + + def cleanup_prefix(self) -> str | None: + return self._cleanup_prefix + + """ Pytest configuration file with global mocks for external dependencies. @@ -6,9 +52,6 @@ during unit tests. The mock is registered in sys.modules before any test imports are executed. """ -import sys -from unittest.mock import MagicMock - # Mock sientia module sientia_mock = MagicMock() sientia_mock.ModelAnalysis = MagicMock diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 5c979e9..4c7b6e6 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -17,9 +17,11 @@ from laborious.activities.storage import Storage @patch('laborious.activities.activities.Gates.__init__') @patch('laborious.activities.activities.ModelMetrics.__init__') @patch('laborious.activities.activities.API.__init__') +@patch('laborious.activities.activities.MinioRepository') @patch('laborious.activities.activities.MetricsController') def test___init__( mock_metrics_controller, + mock_minio_repository, mock_api_init, mock_model_metrics_init, mock_gates_init, @@ -43,6 +45,7 @@ def test___init__( 'secret_key': 'minio123', 'region_name': 'us-east-1', 'default_bucket': 'test', + 'retention_hours': 24, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} @@ -89,7 +92,8 @@ def test___init__( dbname=postgres_config['dbname'], min_connections=postgres_config['min_connections'], max_connections=postgres_config['max_connections'], - minio_config=minio_config, + retention_hours=minio_config['retention_hours'], + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -101,7 +105,7 @@ def test___init__( mlflow_port=mlflow_config['port'], mlflow_username=mlflow_config['username'], mlflow_password=mlflow_config['password'], - minio_config=minio_config, + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -117,6 +121,7 @@ def test___init__( mock_gates_init.assert_called_once_with( ANY, + minio_repository=mock_minio_repository.return_value, logger=logger, notification_handler=notification_handler, metrics_controller=mock_metrics_controller.return_value, @@ -139,6 +144,16 @@ def test___init__( metrics_controller=mock_metrics_controller.return_value, ) + mock_minio_repository.assert_called_once_with( + endpoint_url=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=mock_metrics_controller.return_value, + ) + @mark.asyncio @patch('laborious.activities.activities.Storage') @@ -147,7 +162,9 @@ def test___init__( @patch('laborious.activities.activities.Gates') @patch('laborious.activities.activities.ModelMetrics') @patch('laborious.activities.activities.API') +@patch('laborious.activities.activities.MinioRepository') async def test_shutdown( + _mock_minio_repository, mock_api_init, mock_model_metrics_init, mock_gates_init, @@ -172,6 +189,7 @@ async def test_shutdown( 'secret_key': 'minio123', 'region_name': 'us-east-1', 'default_bucket': 'test', + 'retention_hours': 24, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py index 52c8120..e1e21a9 100644 --- a/tests/laborious/activities/test_api.py +++ b/tests/laborious/activities/test_api.py @@ -61,6 +61,60 @@ def base_input_data(): } +@patch('laborious.activities.api.PIWebAPIClient') +def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client): + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + api_instance = API( + base_url='https://test-pi-server.com', + auth_type='bearer', + auth_token='test_token', + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + with patch.object( + SientiaMonitoring, + 'get_core_labels', + return_value={ + 'pod_id': 'test_pod', + 'model_name': 'test_model', + 'operation_type': '-', + }, + ): + labels = api_instance.get_pi_web_api_core_labels( + metadata=metadata['metadata'], operation_type=None + ) + assert 'operation_type' not in labels + + +@patch('laborious.activities.api.PIWebAPIClient') +def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client): + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + + api_instance = API( + base_url='https://test-pi-server.com', + auth_type='bearer', + auth_token='test_token', + logger=MagicMock(), + notification_handler=MagicMock(), + metrics_controller=AsyncMock(), + ) + with patch.object( + SientiaMonitoring, + 'get_core_labels', + return_value={ + 'pod_id': 'test_pod', + 'model_name': 'test_model', + 'operation_type': 'write', + }, + ): + labels = api_instance.get_pi_web_api_core_labels( + metadata=metadata['metadata'], operation_type='write' + ) + assert labels['operation_type'] == 'write' + + def test__init__(): api = API( base_url='https://test-pi-server.com', diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index e33ea4f..eaee176 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -1,11 +1,29 @@ from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +from pandas import DataFrame from pytest import fixture, mark from sientia_do.notifications.models import NotificationLevel from laborious.activities.gates import Gates +def _minio_payload(retrieve_return, status=None): + """ + Build a MinioDataFramePayload-like test double with async retrieve. + + Args: + retrieve_return: Value returned from await retrieve(minio_repo, metadata). + status: Optional status dict for MLflow response gate (payload.status). + + Return: + MagicMock: Object with async retrieve and optional status. + """ + p = MagicMock() + p.retrieve = AsyncMock(return_value=retrieve_return) + p.status = status + return p + + @fixture def gates_activity(): gates = Gates( @@ -40,7 +58,7 @@ async def test_input_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -65,7 +83,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac input_data = { **metadata, 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': []}, + 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -90,7 +108,7 @@ async def test_input_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -108,7 +126,7 @@ async def test_input_gate_with_filter(gates_activity): input_data = { **metadata, 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': []}, + 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -126,7 +144,7 @@ async def test_input_gate_with_filter_not_caught(gates_activity): input_data = { **metadata, 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -144,7 +162,10 @@ async def test_mlflow_response_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -169,7 +190,10 @@ async def test_mlflow_response_gate_filter_exception( input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -195,7 +219,10 @@ async def test_mlflow_response_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -214,10 +241,10 @@ async def test_mlflow_response_gate_with_filter(gates_activity): input_data = { **metadata, 'filters': {'API_ERROR': {'policy': 'STOP'}}, - 'data': { - 'success': False, - 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, - }, + 'data': _minio_payload( + {'content': {'message': 'API error occurred', 'traceback': 'error trace'}}, + status={'success': False, 'message': 'API error occurred'}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -237,10 +264,10 @@ async def test_mlflow_response_gate_with_filter_not_caught(gates_activity): input_data = { **metadata, 'filters': {'API_ERROR': {'policy': 'STOP'}}, - 'data': { - 'success': True, - 'content': {'message': 'success'}, - }, + 'data': _minio_payload( + {'content': {'message': 'success'}}, + status={'success': True}, + ), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -259,10 +286,7 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity): input_data = { **metadata, 'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}}, - 'data': { - 'success': True, - 'content': {'message': 'success'}, - }, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -287,10 +311,7 @@ async def test_mlflow_content_gate_filter_exception( input_data = { **metadata, 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, - 'data': { - 'success': False, - 'content': {'message': 'API error occurred', 'traceback': 'error trace'}, - }, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -317,7 +338,7 @@ async def test_mlflow_content_gate_no_filters(gates_activity): input_data = { **metadata, 'filters': {}, - 'data': {'value': [1, 2, 3]}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['CONTINUE', 'STOP', 'REPEAT'], } @@ -336,7 +357,7 @@ async def test_mlflow_content_gate_with_filter(gates_activity): input_data = { **metadata, 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, - 'data': {'value': [None, None, None]}, + 'data': _minio_payload(DataFrame({'value': [None, None, None]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -356,7 +377,7 @@ async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): input_data = { **metadata, 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, - 'data': {'content': {'message': 'success'}}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -369,6 +390,22 @@ async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): gates_activity.debug.assert_called() +@mark.asyncio +async def test_mlflow_content_gate_filter_returns_false(gates_activity): + input_data = { + **metadata, + 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, + 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), + 'type': 'test', + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + result = await gates_activity.mlflow_content_gate(input_data) + + assert result == (None, 0, '') + gates_activity.debug.assert_called() + + def test_get_prediction_store_policy_invalid_policy(gates_activity): # Arrange prediction_store_policy = 'INVALID_POLICY' @@ -430,10 +467,14 @@ async def test_format_prediction_no_timestamp(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': {'2023-05-26 11:12:27': 1}, - 'response_time': {'2023-05-26 11:12:27': 0.1}, - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': {'2023-05-26 11:12:27': 1}, + 'response_time': {'2023-05-26 11:12:27': 0.1}, + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:1', @@ -457,18 +498,22 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': { - '2023-05-26 11:12:27': 1, - '2023-05-26 11:12:28': 2, - '2023-05-26 11:12:29': 3, - }, - 'response_time': { - '2023-05-26 11:12:27': 0.1, - '2023-05-26 11:12:28': 0.2, - '2023-05-26 11:12:29': 0.3, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': { + '2023-05-26 11:12:27': 1, + '2023-05-26 11:12:28': 2, + '2023-05-26 11:12:29': 3, + }, + 'response_time': { + '2023-05-26 11:12:27': 0.1, + '2023-05-26 11:12:28': 0.2, + '2023-05-26 11:12:29': 0.3, + }, + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'erl:2', @@ -492,18 +537,22 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': { - '2023-05-26 11:12:27': 1, - '2023-05-26 11:12:28': 2, - '2023-05-26 11:12:29': 3, - }, - 'response_time': { - '2023-05-26 11:12:27': 0.1, - '2023-05-26 11:12:28': 0.2, - '2023-05-26 11:12:29': 0.3, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': { + '2023-05-26 11:12:27': 1, + '2023-05-26 11:12:28': 2, + '2023-05-26 11:12:29': 3, + }, + 'response_time': { + '2023-05-26 11:12:27': 0.1, + '2023-05-26 11:12:28': 0.2, + '2023-05-26 11:12:29': 0.3, + }, + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', @@ -527,11 +576,19 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): # Arrange input_data = { **metadata, - 'data': { - 'prediction': [1, 2, 3], - 'response_time': [0.1, 0.2, 0.3], - 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'], - }, + 'data': _minio_payload( + DataFrame( + { + 'prediction': [1, 2, 3], + 'response_time': [0.1, 0.2, 0.3], + 'timestamp': [ + '2023-05-26 11:12:27', + '2023-05-26 11:12:28', + '2023-05-26 11:12:29', + ], + } + ) + ), 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', @@ -547,76 +604,106 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): @mark.asyncio -async def test_format_transformed_data_single_row(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_single_row(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': { - 'var1': {'2023-05-26 11:12:27': 1.0}, - 'var2': {'2023-05-26 11:12:27': 2.0}, - }, + 'data': _minio_payload( + DataFrame( + { + 'var1': {'2023-05-26 11:12:27': 1.0}, + 'var2': {'2023-05-26 11:12:27': 2.0}, + } + ) + ), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert result['timestamp'] == {0: '2023-05-26 11:12:27', 1: '2023-05-26 11:12:27'} - assert result['variable'] == {0: 'var1', 1: 'var2'} - assert result['value'] == {0: 1.0, 1: 2.0} - assert result['model_id'] == {0: 'test_model', 1: 'test_model'} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @mark.asyncio -async def test_format_transformed_data_multiple_rows(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_multiple_rows(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': { - 'var1': { - '2023-05-26 11:12:27': 1.0, - '2023-05-26 11:12:28': 2.0, - }, - 'var2': { - '2023-05-26 11:12:27': 3.0, - '2023-05-26 11:12:28': 4.0, - }, - }, + 'data': _minio_payload( + DataFrame( + { + 'var1': { + '2023-05-26 11:12:27': 1.0, + '2023-05-26 11:12:28': 2.0, + }, + 'var2': { + '2023-05-26 11:12:27': 3.0, + '2023-05-26 11:12:28': 4.0, + }, + } + ) + ), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert len(result['timestamp']) == 4 - assert len(result['variable']) == 4 - assert len(result['value']) == 4 - assert len(result['model_id']) == 4 - assert all(v == 'test_model' for v in result['model_id'].values()) - assert set(result['variable'].values()) == {'var1', 'var2'} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @mark.asyncio -async def test_format_transformed_data_empty_data(gates_activity): +@patch('laborious.activities.gates.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock) +async def test_format_transformed_data_empty_data(mock_from_dataframe, gates_activity): # Arrange + payload_result = MagicMock() + mock_from_dataframe.return_value = payload_result input_data = { **metadata, - 'data': {}, + 'data': _minio_payload(DataFrame()), 'model_id': 'test_model', + 'model_name': 'test_model', } # Act result = await gates_activity.format_transformed_data(input_data) # Assert - assert result['timestamp'] == {} - assert result['variable'] == {} - assert result['value'] == {} - assert result['model_id'] == {} + assert result is payload_result + mock_from_dataframe.assert_called_once() + kwargs = mock_from_dataframe.call_args.kwargs + assert kwargs['model_name'] == 'test_model' + assert kwargs['operation'] == 'transform' + assert kwargs['workflow_metadata'] == metadata['metadata'] + assert kwargs['minio_repo'] is gates_activity.minio_repository + assert 'dataframe' in kwargs gates_activity.info.assert_called() @@ -711,31 +798,6 @@ async def test_format_retrain_report_failure(gates_activity): gates_activity.debug.assert_called() -@mark.asyncio -async def test_get_last_timestamp_with_data(gates_activity): - # Arrange - input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}} - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert result == '2023-05-26 11:12:28' - - -@mark.asyncio -async def test_get_last_timestamp_no_data(gates_activity): - # Arrange - input_data = {'data': {}, **metadata} - - # Act - result = await gates_activity.get_last_timestamp(input_data) - - # Assert - assert isinstance(result, str) # Should be a timestamp string - assert len(result) > 0 - - @mark.asyncio @patch('laborious.activities.gates.metrics') async def test_write_metrics(mock_metrics, gates_activity): diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 46ae909..baa06e4 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -11,21 +11,27 @@ from laborious.activities.mlflow import MLFlow @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def test___init__(mock_minio_repository, mock_mlflow_repository): + logger = MagicMock() + notification_handler = MagicMock() + metrics_controller = AsyncMock() + minio_repo = mock_minio_repository( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + bucket='test', + ) mlflow = MLFlow( mlflow_host='http://localhost', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', - minio_config={ - 'endpoint_url': 'http://localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), + minio_repository=minio_repo, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, ) assert mlflow.mlflow_host == 'http://localhost' @@ -52,21 +58,27 @@ def test___init__(mock_minio_repository, mock_mlflow_repository): @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def mlflow(mock_minio_repository, mock_mlflow_repository): + logger = MagicMock() + notification_handler = MagicMock() + metrics_controller = AsyncMock() + minio_repo = mock_minio_repository( + endpoint='localhost:9000', + access_key='minio', + secret_key='minio123', + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + bucket='test', + ) mlflow = MLFlow( mlflow_host='http://localhost:5000', mlflow_port=5000, mlflow_username='admin', mlflow_password='admin', - minio_config={ - 'endpoint_url': 'http://localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, - logger=MagicMock(), - notification_handler=MagicMock(), - metrics_controller=AsyncMock(), + minio_repository=minio_repo, + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, ) mlflow.model_monitoring_repository = AsyncMock() @@ -96,165 +108,161 @@ metadata = { @mark.asyncio @patch( - 'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) -@patch('laborious.activities.mlflow.max') -async def test_request_transform_success(mock_max, mock_dataframe_from_wire, mlflow): - mock_max.return_value = '2024-01-02' +async def test_request_transform_success(mock_from_dataframe, mlflow): data_mock = MagicMock() - mock_dataframe_from_wire.return_value = data_mock - # Mock input data + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + input_data = { **metadata, - 'data': [ - { - 'timestamp': '2024-01-01', - 'variable': 'var1', - 'value': 1.0, - 'created_at': '2024-01-01 12:00:00', - }, - { - 'timestamp': '2024-01-01', - 'variable': 'var2', - 'value': 2.0, - 'created_at': '2024-01-01 12:00:00', - }, - { - 'timestamp': '2024-01-02', - 'variable': 'var1', - 'value': 3.0, - 'created_at': '2024-01-02 12:00:00', - }, - { - 'timestamp': '2024-01-02', - 'variable': 'var2', - 'value': 4.0, - 'created_at': '2024-01-02 12:00:00', - }, - { - 'timestamp': '2024-01-02', - 'variable': 'var1', - 'value': 1.0, - 'created_at': '2024-01-01 12:00:00', - }, - { - 'timestamp': '2024-01-02', - 'variable': 'var2', - 'value': 1.0, - 'created_at': '2024-01-01 12:00:00', - }, - ], + 'data': payload, 'model_name': 'test_model', 'model_config': {}, } - # Mock the transform response - expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']} - mlflow.model_monitoring_repository.transform.return_value = expected_response + transform_response = {'success': True, 'content': MagicMock()} + mlflow.model_monitoring_repository.transform.return_value = transform_response data_mock.sort_values.return_value = data_mock data_mock.drop_duplicates.return_value = data_mock data_mock.pivot.return_value = data_mock - # Call the method response_data = await mlflow.request_transform(input_data) - # Verify the data was correctly transformed - data_mock.pivot.assert_called_once_with( - index='timestamp', columns='variable', values='value' - ) - data_mock.fillna.assert_called_once_with(np.nan, inplace=True) - # mock_dataframe.reset_index.assert_called_once() - data_mock.columns.name = None - - # Verify the response - assert response_data == expected_response - - # Verify the repository was called with correct arguments mlflow.model_monitoring_repository.transform.assert_called_once_with( 'test_model', data_mock, {}, metadata['metadata'] ) + mock_from_dataframe.assert_called_once() + assert response_data == mock_from_dataframe.return_value @mark.asyncio @patch( - 'laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', new_callable=AsyncMock, ) -@patch('laborious.activities.mlflow.to_datetime') -@patch('laborious.activities.mlflow.max') -async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe_from_wire, mlflow): - mock_max.return_value = '2024-01-02' +async def test_request_transform_failure(mock_from_dataframe, mlflow): data_mock = MagicMock() - mock_dataframe_from_wire.return_value = data_mock - # Mock input data + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + input_data = { **metadata, - 'data': { - 'variable': { - '2024-01-01': 'var1', - '2024-01-02': 'var2', - '2024-01-03': 'var1', - '2024-01-04': 'var2', - }, - 'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0}, - }, + 'data': payload, 'model_name': 'test_model', 'model_config': {}, } - # Mock the predict response - expected_response = {'prediction': [0.5, 0.6]} - mlflow.model_monitoring_repository.predict.return_value = expected_response + transform_response = {'success': False, 'message': 'Transform failed'} + mlflow.model_monitoring_repository.transform.return_value = transform_response - # Call the method - response_data = await mlflow.request_predict(input_data) + data_mock.sort_values.return_value = data_mock + data_mock.drop_duplicates.return_value = data_mock + data_mock.pivot.return_value = data_mock - data_mock.replace.assert_called_once_with(np.nan, None, inplace=True) - data_mock.__setitem__.assert_any_call( - 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value - ) - data_mock.__setitem__.assert_any_call( - 'timestamp', mock_to_datetime.return_value.dt.strftime.return_value - ) - - mock_to_datetime.assert_called_once_with( - data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ - ) - mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) - - mock_to_datetime.assert_called_once_with( - data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ - ) - mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) - - # Verify the response - assert response_data == expected_response - - # Verify the repository was called with correct arguments - mlflow.model_monitoring_repository.predict.assert_called_once_with( - 'test_model', data_mock, {}, metadata['metadata'] + response_data = await mlflow.request_transform(input_data) + + mock_from_dataframe.assert_called_once_with( + dataframe=None, + minio_repo=mlflow.minio_repository, + model_name='test_model', + operation='transform', + status=transform_response, + workflow_metadata=metadata['metadata'], ) + assert response_data == mock_from_dataframe.return_value @mark.asyncio -@patch('laborious.activities.mlflow.read_parquet') +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mock_read_parquet, mlflow): +async def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + + input_data = { + **metadata, + 'data': payload, + 'model_name': 'test_model', + 'model_config': {}, + } + + predict_response = {'success': True, 'content': MagicMock()} + mlflow.model_monitoring_repository.predict.return_value = predict_response + + response_data = await mlflow.request_predict(input_data) + + data_mock.replace.assert_called_once_with(np.nan, None, inplace=True) + mock_to_datetime.assert_called_once_with( + data_mock.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ + ) + mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT) + + mlflow.model_monitoring_repository.predict.assert_called_once_with( + 'test_model', data_mock, {}, metadata['metadata'] + ) + mock_from_dataframe.assert_called_once() + assert response_data == mock_from_dataframe.return_value + + +@mark.asyncio +@patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, +) +@patch('laborious.activities.mlflow.to_datetime') +async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow): + data_mock = MagicMock() + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=data_mock) + + input_data = { + **metadata, + 'data': payload, + 'model_name': 'test_model', + 'model_config': {}, + } + + predict_response = {'success': False, 'message': 'Predict failed'} + mlflow.model_monitoring_repository.predict.return_value = predict_response + + response_data = await mlflow.request_predict(input_data) + + mock_from_dataframe.assert_called_once_with( + dataframe=None, + minio_repo=mlflow.minio_repository, + model_name='test_model', + operation='predict', + status=predict_response, + workflow_metadata=metadata['metadata'], + ) + assert response_data == mock_from_dataframe.return_value + + +@mark.asyncio +@patch('laborious.activities.mlflow.to_datetime') +async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', } - mlflow.minio_repository.download_file.return_value = b'parquet-bytes' - mock_read_parquet.return_value = MagicMock() + raw_data = MagicMock(columns=['variable', 'timestamp', 'value']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, - 'object_key': 'test_object_key', + 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -264,8 +272,6 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mock } ) - raw_data = mock_read_parquet.return_value - timestamp = raw_data.__getitem__.return_value.max.return_value raw_data.sort_values.assert_not_called() @@ -317,22 +323,22 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mock @mark.asyncio -@patch('laborious.activities.mlflow.MinioDataFramePayload.dataframe_from_wire', new_callable=AsyncMock) @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_with_payload_data( - mock_to_datetime, mock_dataframe_from_wire, mlflow -): +async def test_retrain_model_success_with_payload_data(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': True, 'experiment': 'test_experiment', 'message': 'Model retrained successfully.', } - mock_dataframe_from_wire.return_value = MagicMock() + + raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, - 'data': {'data': {'a': [1]}}, + 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -347,24 +353,22 @@ async def test_retrain_model_success_with_payload_data( @mark.asyncio -@patch('laborious.activities.mlflow.read_parquet') @patch('laborious.activities.mlflow.to_datetime') -async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mock_read_parquet, mlflow): +async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow): mlflow.model_monitoring_repository.retrain_model.return_value = { 'success': False, 'traceback': 'test_traceback', 'message': 'Model retrained failed.', } - mlflow.minio_repository.download_file.return_value = b'parquet-bytes' - mock_read_parquet.return_value = MagicMock( - columns=['variable', 'timestamp', 'value', 'created_at'] - ) + raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at']) + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=raw_data) response = await mlflow.retrain_model( { **metadata, - 'object_key': 'test_object_key', + 'data': payload, 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -374,8 +378,6 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mock_re } ) - raw_data = mock_read_parquet.return_value - timestamp = raw_data.__getitem__.return_value.max.return_value raw_data.sort_values.assert_called_once_with('created_at', ascending=False) @@ -439,14 +441,9 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mock_re @mark.asyncio async def test_retrain_model_data_error(mlflow): - mlflow.minio_repository.download_file.side_effect = Exception( - 'Error loading retrain data' - ) - response = await mlflow.retrain_model( { **metadata, - 'object_key': 'test_object_key', 'model_name': 'test_model', 'model_config': { 'target': 'target', @@ -458,7 +455,7 @@ async def test_retrain_model_data_error(mlflow): assert response == { 'success': False, - 'message': 'Error loading retrain data: Error loading retrain data', + 'message': "Error loading retrain data: 'data'", 'traceback': ANY, 'timestamp': ANY, } diff --git a/tests/laborious/activities/test_model_metrics.py b/tests/laborious/activities/test_model_metrics.py index f20687d..b2b5e0b 100644 --- a/tests/laborious/activities/test_model_metrics.py +++ b/tests/laborious/activities/test_model_metrics.py @@ -743,6 +743,101 @@ async def test_get_drift_metrics_univariate_error( raise AssertionError('Expected Exception') +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_multivariate_error( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + mock_time.return_value = 1000.0 + + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.detect_multivariate_drift.side_effect = Exception( + 'Multivariate drift error' + ) + + reference_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + target_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + try: + await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + except Exception as e: + assert str(e) == 'Multivariate drift error' + model_metrics_activity.error.assert_called_once_with( + 'Error detecting multivariate drift: Multivariate drift error', metadata['metadata'] + ) + model_metrics_activity.emit_metric.assert_called_with( + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY + ) + else: + raise AssertionError('Expected Exception') + + +@mark.asyncio +@patch('laborious.activities.model_metrics.to_datetime') +@patch('laborious.activities.model_metrics.time.time') +@patch('laborious.activities.model_metrics.ModelAnalysis') +@patch('laborious.activities.model_metrics.metrics') +async def test_get_drift_metrics_dataframe_error( + mock_metrics, mock_model_analysis, mock_time, mock_to_datetime, model_metrics_activity +): + mock_time.return_value = 1000.0 + + mock_model_analysis.return_value.detect_univariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.detect_multivariate_drift.return_value = MagicMock() + mock_model_analysis.return_value.get_drift_metrics_dataframe.side_effect = Exception( + 'Dataframe error' + ) + + reference_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + target_data = DataFrame( + {'timestamp': ['2023-05-26 11:12:27'], 'target': [1.0], 'feature1': [1.0]} + ) + reference_columns = reference_data.drop( + columns=['target', 'timestamp'], errors='ignore' + ).columns + + try: + await model_metrics_activity.get_drift_metrics( + reference_data=reference_data, + target_data=target_data, + target_name='target', + reference_columns=reference_columns, + drift_metrics=['ks_test'], + chunk_period='min', + metadata=metadata['metadata'], + ) + except Exception as e: + assert str(e) == 'Dataframe error' + model_metrics_activity.error.assert_called_once_with( + 'Error getting drift metrics: Dataframe error', metadata['metadata'] + ) + model_metrics_activity.emit_metric.assert_called_with( + metric_object=mock_metrics.MODEL_ANALYZE_ERROR_COUNT, tags=ANY + ) + else: + raise AssertionError('Expected Exception') + + @mark.asyncio async def test_calculate_simple_metrics_success_all_metrics(model_metrics_activity): # Arrange @@ -987,3 +1082,27 @@ async def test_calculate_simple_metrics_success_multiple_metrics_subset(model_me model_metrics_activity.info.assert_called_once_with( "Calculating simple metrics for model test_model_id: ['rmse', 'mae']", metadata['metadata'] ) + + +@mark.asyncio +async def test_calculate_simple_metrics_unknown_metric_ignored(model_metrics_activity): + target_data = DataFrame( + { + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28'], + 'target': [1.0, 2.0], + 'prediction': [1.1, 2.1], + } + ) + + input_data = { + **metadata, + 'model_id': 'test_model_id', + 'target_data': target_data.to_dict(), + 'metrics': ['unknown_metric', 'rmse'], + 'interval_minutes': 5, + } + + result = DataFrame(await model_metrics_activity.calculate_simple_metrics(input_data)) + + assert len(result['metric']) == 1 + assert result['metric'].values[0] == 'rmse' diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index 963dbf2..f0272d1 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -29,13 +29,8 @@ def storage(mock_minio_repository): dbname='postgres', min_connections=1, max_connections=10, - minio_config={ - 'endpoint_url': 'localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, + retention_hours=24, + minio_repository=mock_minio_repository.return_value, logger=MagicMock(), notification_handler=MagicMock(), metrics_controller=AsyncMock(), @@ -47,6 +42,7 @@ def test___init___not_hasattr(mock_minio_repository): logger = MagicMock() notification_handler = MagicMock() metrics_controller = AsyncMock() + minio_repo = mock_minio_repository.return_value storage = Storage( host='localhost', port=5432, @@ -55,28 +51,16 @@ def test___init___not_hasattr(mock_minio_repository): dbname='postgres', min_connections=1, max_connections=10, - minio_config={ - 'endpoint_url': 'localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, + retention_hours=24, + minio_repository=minio_repo, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) assert isinstance(storage, Postgres) - mock_minio_repository.assert_called_once_with( - endpoint='localhost:9000', - access_key='minio', - secret_key='minio123', - logger=logger, - notification_handler=notification_handler, - metrics_controller=metrics_controller, - bucket='test', - ) + assert storage.minio_repository is minio_repo + mock_minio_repository.assert_not_called() @patch('laborious.activities.storage.MinioRepository') @@ -93,27 +77,15 @@ def test___init___none_minio_repository(mock_minio_repository, storage): dbname='postgres', min_connections=1, max_connections=10, - minio_config={ - 'endpoint_url': 'localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, + retention_hours=24, + minio_repository=None, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, ) - mock_minio_repository.assert_called_once_with( - endpoint='localhost:9000', - access_key='minio', - secret_key='minio123', - logger=logger, - notification_handler=notification_handler, - metrics_controller=metrics_controller, - bucket='test', - ) + assert storage.minio_repository is None + mock_minio_repository.assert_not_called() @patch('laborious.activities.storage.MinioRepository') @@ -126,13 +98,8 @@ def test___init___done_repository(mock_minio_repository, storage): dbname='postgres', min_connections=1, max_connections=10, - minio_config={ - 'endpoint_url': 'localhost:9000', - 'access_key': 'minio', - 'secret_key': 'minio123', - 'region_name': 'us-east-1', - 'default_bucket': 'test', - }, + retention_hours=24, + minio_repository=mock_minio_repository.return_value, logger=MagicMock(), notification_handler=MagicMock(), metrics_controller=AsyncMock(), @@ -231,55 +198,58 @@ def test___del__(storage): storage.close.assert_called_once() -def test_estimate_payload_size_bytes(storage): - assert storage._estimate_payload_size_bytes({'x': 1}) > 0 - - @mark.asyncio async def test_load_query_with_minio_offload_no_rows(storage): storage.load_custom_query = AsyncMock(return_value=None) - result = await storage.load_query_with_minio_offload( - {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} - ) - assert result['success'] is False + storage_result = {'success': False} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: + result = await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} + ) + assert result == storage_result + mock_from_dataframe.assert_awaited_once() @mark.asyncio async def test_load_query_with_minio_offload_inline(storage): storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) - result = await storage.load_query_with_minio_offload( - {**metadata, 'query': 'SELECT 1', 'model_name': 'my-model', 'key_prefix': 'predictions/s'} - ) - assert result.get('success') is True - assert 'data' in result - assert result.get('object_key') is None + storage_result = {'success': True, 'data': {'a': [1]}, 'object_key': None} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: + result = await storage.load_query_with_minio_offload( + { + **metadata, + 'query': 'SELECT 1', + 'model_name': 'my-model', + 'key_prefix': 'predictions/s', + } + ) + assert result == storage_result + mock_from_dataframe.assert_awaited_once() @mark.asyncio -@patch('laborious.utils.models.minio_dataframe_payload.MinioDataFramePayload.estimate_size_bytes') -async def test_load_query_with_minio_offload_minio(mock_estimate, storage): - mock_estimate.return_value = 10**9 +async def test_load_query_with_minio_offload_minio(storage): storage.load_custom_query = AsyncMock(return_value=[{'a': 1}]) - storage.minio_repository.upload_file = AsyncMock( - return_value={ - 'minio_object_name': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet' - } - ) - storage.minio_repository.bucket = 'test' - - fixed = datetime.datetime(2024, 1, 15, 12, 30, 45) - with patch('laborious.utils.models.minio_dataframe_payload.now', return_value=fixed): + storage_result = {'success': True, 'data': None, 'object_key': 'object-key'} + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dataframe', + new_callable=AsyncMock, + return_value=storage_result, + ) as mock_from_dataframe: result = await storage.load_query_with_minio_offload( {**metadata, 'query': 'SELECT 1', 'model_name': 'm', 'key_prefix': 'predictions/s'} ) - assert result.get('success') is True - assert result.get('data') is None - assert ( - result['object_key'] - == 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-15_12-30-45.parquet' - ) - storage.minio_repository.upload_file.assert_called_once() + assert result == storage_result + mock_from_dataframe.assert_awaited_once() @mark.asyncio @@ -297,12 +267,117 @@ async def test_cleanup_minio_objects_expired(mock_now, storage): storage.send_notification_async = AsyncMock() result = await storage.cleanup_minio_objects_expired( - {**metadata, 'prefixes': ['training_datasets/m']} + {**metadata, 'prefix': 'training_datasets/m'} ) - assert result['success'] is True assert result['deleted_count'] == 1 + assert result['failed_count'] == 0 + deleted_key = ( + 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet' + ) + assert deleted_key in result['deleted'] + assert result['deleted'][deleted_key]['success'] is True + storage.minio_repository.list_objects.assert_called_once_with( + prefix='training_datasets/m', + recursive=True, + metadata=metadata['metadata'], + ) storage.minio_repository.delete_file.assert_called_once_with( object_name='sientia/streamlit-connectors/training_datasets/m/m-initial-2024-12-01_00-00-00.parquet', metadata=metadata['metadata'], ) + + +@mark.asyncio +async def test_load_query_with_minio_offload_minio_not_initialized(storage): + storage.minio_repository = None + + with raises(ValueError, match='Minio repository not initialized'): + await storage.load_query_with_minio_offload( + {**metadata, 'query': 'SELECT 1', 'model_name': 'm'} + ) + + +@mark.asyncio +async def test_export_payload_to_postgres(storage): + payload = AsyncMock() + payload.retrieve = AsyncMock(return_value=MagicMock()) + storage.export_data_to_postgres = AsyncMock(return_value={'success': True}) + + result = await storage.export_payload_to_postgres( + {**metadata, 'data': payload, 'schema': 'public', 'table': 't'} + ) + + payload.retrieve.assert_awaited_once_with(storage.minio_repository, metadata['metadata']) + storage.export_data_to_postgres.assert_awaited_once() + assert result == {'success': True} + + +@mark.asyncio +async def test_cleanup_minio_objects_expired_minio_not_initialized(storage): + storage.minio_repository = None + + with raises(ValueError, match='Minio repository not initialized'): + await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'}) + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock( + return_value=['some/random/key-without-timestamp.parquet'] + ) + storage.minio_repository.delete_file = AsyncMock() + storage.send_notification_async = AsyncMock() + + result = await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'}) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 0 + storage.minio_repository.delete_file.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + old_key = 'training_datasets/m/m-initial-2024-12-01_00-00-00.parquet' + storage.minio_repository.list_objects = AsyncMock(return_value=[old_key]) + storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error')) + storage.send_notification_async = AsyncMock() + + result = await storage.cleanup_minio_objects_expired( + {**metadata, 'prefix': 'training_datasets/m'} + ) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 1 + assert old_key in result['failed'] + assert result['failed'][old_key]['success'] is False + assert result['failed'][old_key]['message'] == 'delete error' + + +@mark.asyncio +@patch('laborious.activities.storage.now') +async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storage): + mock_now.return_value = datetime.datetime(2025, 1, 10, 12, 0, 0) + storage.minio_repository.list_objects = AsyncMock(side_effect=Exception('list error')) + storage.send_notification_async = AsyncMock() + storage.error = MagicMock() + + result = await storage.cleanup_minio_objects_expired( + {**metadata, 'prefix': 'training_datasets/m'} + ) + + assert result['deleted_count'] == 0 + assert result['failed_count'] == 0 + storage.send_notification_async.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED', + message='Error cleaning up MinIO objects: list error', + block='cleanup_minio_objects_expired', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + storage.error.assert_called_once() diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py index 71aeda0..9349264 100644 --- a/tests/laborious/utils/models/test_minio_dataframe_payload.py +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -1,8 +1,14 @@ from datetime import datetime +from io import BytesIO +from unittest.mock import AsyncMock, MagicMock, patch -from pytest import mark +import pytest +from pandas import DataFrame -from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload +from laborious.utils.models.minio_dataframe_payload import ( + MinioDataFramePayload, + _build_object_key, +) def test_parse_object_timestamp_hyphenated_model(): @@ -21,34 +27,190 @@ def test_parse_object_timestamp_invalid(): assert MinioDataFramePayload.parse_object_timestamp('bad.parquet') is None -def test_is_offloaded_dict_true_false(): - assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': None}) is True - assert MinioDataFramePayload.is_offloaded_dict({'object_key': 'k', 'data': {}}) is False - assert MinioDataFramePayload.is_offloaded_dict({'data': {}}) is False +def test_estimate_size_bytes_returns_positive_for_nonempty_frame(): + df = DataFrame({'a': [1, 2]}) + size = MinioDataFramePayload.estimate_size_bytes(df) + assert isinstance(size, int) + assert size > 0 -def test_cleanup_prefix_from_payload_dict(): - p = { - 'object_key': 'sientia/streamlit-connectors/training_datasets/m/m-initial-2024-01-01_00-00-00.parquet', - 'bucket': 'b', - 'data': None, - } - assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'training_datasets/m' +def test_cleanup_prefix_when_offloaded_returns_object_prefix(): + payload = MinioDataFramePayload( + last_timestamp='t', + data=None, + object_key='training_datasets/m/m-initial-2024-01-01_00-00-00.parquet', + object_prefix='training_datasets/m', + ) + assert MinioDataFramePayload.cleanup_prefix(payload) == 'training_datasets/m' -def test_cleanup_prefix_from_explicit_object_prefix(): - p = {'object_key': 'x.parquet', 'object_prefix': 'my/prefix', 'data': None} - assert MinioDataFramePayload.cleanup_prefix_from_payload_dict(p) == 'my/prefix' +def test_cleanup_prefix_when_inline_returns_none(): + payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None) + assert MinioDataFramePayload.cleanup_prefix(payload) is None -@mark.asyncio -async def test_resolve_dict_if_offloaded_noop(): - d = {'success': True, 'data': {'a': [1]}} - out = await MinioDataFramePayload.resolve_dict_if_offloaded(d, None, {}) - assert out is d +def test_has_data_true_when_object_key_set(): + payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key='k') + assert payload.has_data() is True -@mark.asyncio -async def test_dataframe_from_wire_list(): - df = await MinioDataFramePayload.dataframe_from_wire([{'a': 1}], None, {}) - assert list(df.columns) == ['a'] +@pytest.mark.asyncio +async def test_retrieve_inline_dict_as_dataframe(): + payload = MinioDataFramePayload(last_timestamp='t', data={'a': [1, 2]}) + minio = AsyncMock() + out = await payload.retrieve(minio, {'metadata': {}}) + assert list(out.columns) == ['a'] + minio.download_file.assert_not_called() + + +@pytest.mark.asyncio +async def test_retrieve_downloads_parquet_when_offloaded(): + source = DataFrame({'a': [1, 2]}) + buf = BytesIO() + source.to_parquet(buf, engine='pyarrow', index=True) + file_bytes = buf.getvalue() + + payload = MinioDataFramePayload( + last_timestamp='t', + data=None, + object_key='training_datasets/m/f.parquet', + object_prefix='training_datasets/m', + ) + minio = AsyncMock() + minio.download_file = AsyncMock(return_value=file_bytes) + + out = await payload.retrieve(minio, {'metadata': {}}) + + minio.download_file.assert_awaited_once_with( + object_name='training_datasets/m/f.parquet', + metadata={'metadata': {}}, + ) + assert list(out.columns) == ['a'] + + +def test_build_object_key(): + key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00') + assert key == 'training_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet' + assert prefix == 'training_datasets/my-model' + + +def test_build_object_key_strips_slashes(): + key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45') + assert prefix == 'training_datasets/my-model' + assert key.startswith('training_datasets/my-model/') + + +def test_estimate_size_bytes_fallback(): + df = DataFrame({'a': [1, 2]}) + original_to_dict = df.to_dict + df.to_dict = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('to_dict failed')) + size = MinioDataFramePayload.estimate_size_bytes(df) + df.to_dict = original_to_dict + assert isinstance(size, int) + assert size > 0 + + +def test_parse_object_timestamp_bad_datetime(): + key = 'p/m-initial-9999-99-99_99-99-99.parquet' + assert MinioDataFramePayload.parse_object_timestamp(key) is None + + +@pytest.mark.asyncio +async def test_retrieve_empty_when_no_data(): + payload = MinioDataFramePayload(last_timestamp='t', data=None, object_key=None) + minio = AsyncMock() + out = await payload.retrieve(minio, {}) + assert out.empty + minio.download_file.assert_not_called() + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +async def test_from_dataframe_none(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + result = await MinioDataFramePayload.from_dataframe( + dataframe=None, + minio_repo=minio, + model_name='m', + operation='initial', + status={'success': False, 'message': 'no data'}, + ) + assert result.data is None + assert result.status == {'success': False, 'message': 'no data'} + assert result.object_key is None + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +async def test_from_dataframe_empty(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + mock_df = MagicMock() + mock_df.__bool__ = MagicMock(return_value=True) + mock_df.empty = True + result = await MinioDataFramePayload.from_dataframe( + dataframe=mock_df, + minio_repo=minio, + model_name='m', + operation='initial', + ) + assert result.data is None + assert result.object_key is None + + +def _mock_dataframe(data_dict, timestamp_values=None): + """Build a MagicMock that behaves enough like a DataFrame for from_dataframe.""" + mock_df = MagicMock() + mock_df.__bool__ = MagicMock(return_value=True) + mock_df.empty = False + if timestamp_values is None: + timestamp_values = data_dict.get('timestamp', ['2024-01-01']) + ts_col = MagicMock() + ts_col.values.tolist.return_value = timestamp_values + mock_df.__getitem__ = MagicMock(return_value=ts_col) + mock_df.to_dict.return_value = data_dict + buf = BytesIO() + DataFrame(data_dict).to_parquet(buf, engine='pyarrow', index=True) + mock_df.to_parquet = MagicMock(side_effect=lambda b, **kw: b.write(buf.getvalue())) + return mock_df + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 10**9) +async def test_from_dataframe_inline(): + minio = AsyncMock() + df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]}) + result = await MinioDataFramePayload.from_dataframe( + dataframe=df, + minio_repo=minio, + model_name='m', + operation='initial', + ) + assert result.data is not None + assert result.object_key is None + assert result.last_timestamp == '2024-01-01' + + +@pytest.mark.asyncio +@patch('laborious.utils.models.minio_dataframe_payload.now') +@patch('laborious.utils.models.minio_dataframe_payload.OFFLOAD_THRESHOLD_BYTES', 0) +async def test_from_dataframe_offloaded(mock_now): + mock_now.return_value = datetime(2024, 1, 1, 0, 0, 0) + minio = AsyncMock() + minio.upload_file = AsyncMock(return_value={'minio_object_name': 'full/key.parquet'}) + minio.bucket = 'test-bucket' + + df = _mock_dataframe({'timestamp': ['2024-01-01'], 'value': [42]}) + result = await MinioDataFramePayload.from_dataframe( + dataframe=df, + minio_repo=minio, + model_name='m', + operation='initial', + workflow_metadata={'wf': 'data'}, + ) + assert result.data is None + assert result.object_key == 'full/key.parquet' + assert result.bucket == 'test-bucket' + assert result.uri == 's3://test-bucket/full/key.parquet' + minio.upload_file.assert_awaited_once() diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 5311a22..15f8c6b 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -2,6 +2,7 @@ from datetime import UTC, datetime from unittest.mock import ANY, AsyncMock, MagicMock, call, patch import mlflow as mlflow_lib +import numpy as np import pytest from pandas import DataFrame, Timestamp @@ -1310,10 +1311,8 @@ async def test_transform_success(mlflow_repository): mlflow_repository.get_cached_operation.return_value, metadata['metadata'] ) - assert output == { - 'success': True, - 'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.return_value, - } + assert output['success'] is True + assert output['content'] is mlflow_repository.detect_and_parse_datetime_index.return_value @pytest.mark.asyncio @@ -1355,7 +1354,7 @@ async def test_predict_success_array(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1363,10 +1362,12 @@ async def test_predict_success_array(mlflow_repository): ) assert output['success'] is True - assert output['content'] == { - 'prediction': {'index_1': 2, 'index_2': 3}, - 'response_time': {'index_1': ANY, 'index_2': ANY}, - } + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert list(content.columns) == ['prediction', 'response_time'] + assert content.index.tolist() == data.index.tolist() @pytest.mark.asyncio @@ -1383,7 +1384,7 @@ async def test_predict_success_df(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1391,10 +1392,12 @@ async def test_predict_success_df(mlflow_repository): ) assert output['success'] is True - assert output['content'] == { - 'prediction': {'index_1': 2, 'index_2': 3}, - 'response_time': {'index_1': ANY, 'index_2': ANY}, - } + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert list(content.columns) == ['prediction', 'response_time'] + assert content.index.tolist() == data.index.tolist() @pytest.mark.asyncio @@ -1409,7 +1412,7 @@ async def test_predict_error(mlflow_repository): mlflow_repository.get_cached_operation.assert_called_once_with( model_name=model_name, - data=data, + data=ANY, operation='predict', retention=60, flavor='pyfunc', @@ -1559,3 +1562,119 @@ def test_get_prediction_data_pyfunc(mlflow_repository): assert 'target' in result.columns assert 'timestamp' in result.columns assert result.index.tolist() == [0, 1] + + +@patch('laborious.utils.repository.model_repository.pd.merge') +@patch('laborious.utils.repository.model_repository.isinstance') +@pytest.mark.asyncio +async def test_fit_models_skip_transform(isinstance_mock, pd_merge, mlflow_repository): + isinstance_mock.return_value = True + + data_model = MagicMock(target_variable='feat_2') + prediction_model = MagicMock() + mlflow_repository.download_model = AsyncMock( + side_effect=[(data_model, 'artifact_path'), (prediction_model, 'artifact_path')], + ) + mlflow_repository.detect_and_parse_datetime_index = MagicMock( + return_value=MagicMock( + drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1'])) + ) + ) + mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame()) + + data = MagicMock() + + output = await mlflow_repository.fit_models( + 'model_name', + data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + True, + 'pyfunc', + 'feat_1', + ) + + data_model.fit.assert_not_called() + + assert output['data_model'] == {'model': data_model, 'artifact_path': 'artifact_path'} + + +@patch('laborious.utils.repository.model_repository.force_memory_release') +@patch('laborious.utils.repository.model_repository.path') +@patch('laborious.utils.repository.model_repository.rmtree') +@pytest.mark.asyncio +async def test_create_new_experiment_path_not_exists( + _rmtree, path, force_memory_release, mlflow, mlflow_repository +): + model_name = 'model_name' + data = MagicMock() + prediction_data = MagicMock(spec=DataFrame) + retrain_data = { + 'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'}, + 'prediction_data': prediction_data, + } + + mlflow_repository.get_model_params = MagicMock( + return_value={ + 'transform_flavor': 'sklearn', + 'predict_flavor': 'pyfunc', + 'target_name': 'target_name', + } + ) + + mlflow_repository.get_experiment = MagicMock() + mlflow_repository.get_next_run_name = MagicMock() + mlflow_repository.log_model = AsyncMock() + path.exists.return_value = False + path.join.return_value = './tmp/artifacts/model_name' + + await mlflow_repository.create_new_experiment( + model_name, + data, + retrain_data, + 'latest_production_id', + metadata['metadata'], + 'sklearn', + 'pyfunc', + ) + + _rmtree.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_production_model_by_run_id_transition_error(mlflow, mlflow_repository): + mlflow_repository.client.get_registered_model.return_value = MagicMock( + latest_versions=[ + MagicMock(version='1'), + MagicMock(version='2'), + ] + ) + mlflow_repository.client.transition_model_version_stage.side_effect = Exception( + 'transition error' + ) + + with pytest.raises(Exception, match='transition error'): + await mlflow_repository.update_production_model_by_run_id('0', 'test', metadata['metadata']) + + mlflow_repository.emit_metric.assert_called_with( + metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=ANY + ) + + +@pytest.mark.asyncio +async def test_predict_success_ndarray(mlflow_repository): + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + model_config = {'retention_minutes': 60, 'predict_flavor': 'pyfunc'} + model_name = 'model' + mlflow_repository.get_cached_operation = AsyncMock(return_value=np.array([5.0, 6.0])) + + output = await mlflow_repository.predict(model_name, data, model_config, metadata['metadata']) + + assert output['success'] is True + content = output['content'] + assert isinstance(content, DataFrame) + assert 'prediction' in content.columns + assert 'response_time' in content.columns + assert content.index.tolist() == data.index.tolist() diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 80eb19f..e3ba910 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -102,6 +102,7 @@ def test_build_minio_config_with_env_vars(): 'secret_key': 'test-secret', 'region_name': 'test-region', 'default_bucket': 'test-bucket', + 'retention_hours': 24, } @@ -117,4 +118,5 @@ def test_build_minio_config_with_defaults(): 'secret_key': 'minioadmin', 'region_name': 'us-east-1', 'default_bucket': 'laborious', + 'retention_hours': 24, } diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index dd95ee9..6e8817f 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -34,6 +34,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -58,12 +59,13 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): call( Activities.format_prediction, { + **metadata, 'data': input_data['data'], 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], 'prediction_store_policy': input_data['prediction_store_policy'], - **metadata, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -141,6 +143,7 @@ async def test_run_none_path_flag_with_transformed_data( 'transformed_data': {'transformed': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0.9, 'schema': 'test_schema', 'table_name': 'test_table', @@ -176,12 +179,13 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.format_prediction, { + **metadata, 'data': input_data['data'], 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], 'prediction_store_policy': input_data['prediction_store_policy'], - **metadata, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -189,9 +193,10 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.format_transformed_data, { + **metadata, 'data': input_data['transformed_data'], 'model_id': input_data['model_id'], - **metadata, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -201,8 +206,9 @@ async def test_run_none_path_flag_with_transformed_data( # Assert - start_activity_method for transformed data export workflow_mock.start_activity_method.assert_called_once_with( - Activities.export_data_to_postgres, + Activities.export_payload_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['transform_table_name'], 'data': transformed_data, @@ -210,7 +216,6 @@ async def test_run_none_path_flag_with_transformed_data( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -287,6 +292,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -311,11 +317,11 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction call( Activities.format_default_prediction, { + **metadata, 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], 'comment': input_data['comment'], - **metadata, }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -389,6 +395,7 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -415,12 +422,13 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo call( Activities.format_prediction, { + **metadata, 'data': input_data['data'], 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], 'prediction_store_policy': input_data['prediction_store_policy'], - **metadata, + 'model_name': input_data['model_name'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -496,6 +504,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', @@ -597,6 +606,7 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e 'data': {'test': 'data'}, 'timestamp': '2021-01-01', 'model_id': 1, + 'model_name': metadata['metadata']['model_name'], 'prediction_confidence': 0, 'schema': 'test_schema', 'table_name': 'test_table', diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 4132537..0da262b 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark @@ -27,9 +27,12 @@ metadata = { async def test_run(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=False) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -47,7 +50,6 @@ async def test_run(workflow_mock, prediction_process): # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) @@ -63,21 +65,7 @@ async def test_run(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - **metadata, - 'data': input_data['data'], - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 6 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -102,7 +90,6 @@ async def test_run(workflow_mock, prediction_process): 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], - 'key_prefix': 'predictions/test_schedule', }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -201,9 +188,12 @@ async def test_run(workflow_mock, prediction_process): async def test_run_stop_at_input_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(return_value=True) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -219,7 +209,6 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('stop', 0.95, 'Input data with bad quality'), # input_gate ] @@ -227,18 +216,9 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 2 + assert workflow_mock.execute_local_activity_method.call_count == 1 workflow_mock.execute_local_activity_method.assert_has_calls( [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ), call( Activities.input_gate, { @@ -260,9 +240,12 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, True]) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -278,7 +261,6 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('repeat', 0.95, 'Input data with bad quality'), # input_gate {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data ('continue', 0.95, 'Error'), # mlflow_response_gate (transform) @@ -288,20 +270,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 4 - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 3 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -325,7 +294,6 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], - 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, @@ -357,9 +325,12 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, True]) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -375,7 +346,6 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) @@ -388,21 +358,8 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 5 + assert workflow_mock.execute_local_activity_method.call_count == 4 - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -426,7 +383,6 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], - 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, @@ -474,9 +430,12 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_process): prediction_process.path_flag_handler = AsyncMock(side_effect=[False, False, False, True]) # Arrange + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' input_data = { 'metadata': metadata, - 'data': {'test': 'data'}, + 'data': data_payload, 'schema': 'test_schema', 'table_name': 'test_table', 'transform_table_name': 'test_transform_table', @@ -492,7 +451,6 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p # Mock the activity responses workflow_mock.execute_local_activity_method.side_effect = [ - '2024-01-01', # get_last_timestamp ('continue', 0.95, 'Input data with bad quality'), # input_gate {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) @@ -507,20 +465,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 7 - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.get_last_timestamp, - { - 'data': input_data['data'], - **metadata, - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ) - ] - ) + assert workflow_mock.execute_local_activity_method.call_count == 6 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -544,7 +489,6 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p 'data': input_data['data'], 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], - 'key_prefix': 'predictions/test_schedule', **metadata, }, retry_policy=ANY, @@ -821,3 +765,49 @@ async def test_path_flag_handler_unknown(workflow_mock, prediction_process): assert result is False workflow_mock.execute_activity_method.assert_not_called() workflow_mock.execute_child_workflow.assert_not_called() + + +@mark.asyncio +@patch('laborious.workflows.sub_workflows.prediction_process.workflow', new_callable=AsyncMock) +async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): + prediction_process.path_flag_handler = AsyncMock(return_value=False) + prediction_process.cleanup_prefixes = {'training_datasets/test'} + + data_payload = MagicMock() + data_payload.cleanup_prefix.return_value = 'training_datasets/test' + data_payload.last_timestamp = '2024-01-01' + input_data = { + 'metadata': metadata, + 'data': data_payload, + 'schema': 'test_schema', + 'table_name': 'test_table', + 'transform_table_name': 'test_transform_table', + 'model_id': 1, + 'input_filters': {'test': 'filter'}, + 'mlflow_transform_filters': {'test': 'filter'}, + 'mlflow_predict_filters': {'test': 'filter'}, + 'model_name': 'test_model_name', + 'model_config': {'retention': '30'}, + 'path_priority': ['continue', 'repeat', 'stop'], + 'opc_output_config': {'test': 'config'}, + 'pi_web_api_output_config': {'test': 'config'}, + 'prediction_store_policy': 'lts:1', + } + + workflow_mock.execute_local_activity_method.side_effect = [ + ('continue', 0.95, 'ok'), + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + ('continue', 0.95, ''), + ('continue', 0.95, ''), + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + ('continue', 0.95, ''), + ] + + await prediction_process.run(input_data) + + workflow_mock.execute_activity_method.assert_any_call( + Activities.cleanup_minio_objects_expired, + {**metadata, 'prefix': 'training_datasets/test'}, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index 6413019..be49c89 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark @@ -39,9 +39,12 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): }, } + storage_result = MagicMock() + storage_result.has_data.return_value = True + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'data': {'a': [1]}, 'success': True}, + storage_result, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -77,7 +80,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): Activities.retrain_model, { **metadata, - 'data': {'data': {'a': [1]}, 'success': True}, + 'data': storage_result, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, @@ -160,9 +163,12 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim }, } + storage_result = MagicMock() + storage_result.has_data.return_value = False + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'success': False, 'message': 'No data returned from query'}, + storage_result, {'success': True, 'experiment': 'test_experiment'}, { 'success': True, @@ -174,7 +180,10 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim ] ) - await minimal_retrain.run(input_data) + from pytest import raises + + with raises(ValueError, match='No data returned from query'): + await minimal_retrain.run(input_data) workflow_mock.execute_activity_method.assert_called_once_with( Activities.load_query_with_minio_offload, @@ -209,9 +218,12 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim }, } + storage_result = MagicMock() + storage_result.has_data.return_value = True + workflow_mock.execute_activity_method = AsyncMock( side_effect=[ - {'data': {'a': [1]}, 'success': True}, + storage_result, {'success': False, 'experiment': 'test_experiment'}, { 'success': True, @@ -247,7 +259,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim Activities.retrain_model, { **metadata, - 'data': {'data': {'a': [1]}, 'success': True}, + 'data': storage_result, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 5505e89..8ebd040 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -54,7 +54,6 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': input_data['model_name'], - 'key_prefix': f"predictions/{input_data['schedule_name']}", }, retry_policy=ANY, start_to_close_timeout=ANY, diff --git a/validate.sh b/validate.sh deleted file mode 100755 index 6c72694..0000000 --- a/validate.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/bin/bash -# Model Manager Code Validation Script -# This script runs all code quality checks before committing or deploying - -set -e # Exit on any error - -# Colors for output -RED='\033[0;31m' -GREEN='\033[0;32m' -YELLOW='\033[1;33m' -BLUE='\033[0;34m' -NC='\033[0m' # No Color - -# Args -FIX_MODE=false -while [[ $# -gt 0 ]]; do - case "$1" in - --fix) - FIX_MODE=true - shift - ;; - -h|--help) - echo "Usage: $0 [--fix]" - echo " --fix Apply Ruff auto-fixes (format and lint fixes)." - exit 0 - ;; - *) - echo -e "${RED}Unknown option: $1${NC}" - echo "Usage: $0 [--fix]" - exit 2 - ;; - esac -done - -echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" -echo -e "${BLUE}║ Model Manager - Code Validation Suite ║${NC}" -echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}" -echo "" - -# Check if virtual environment is activated -if [[ -z "${VIRTUAL_ENV}" ]] && [[ -z "${CONDA_DEFAULT_ENV}" ]]; then - echo -e "${YELLOW}⚠️ Warning: No virtual environment detected${NC}" - echo -e "${YELLOW} Consider activating your venv/conda environment${NC}" - echo "" -fi - -# Function to run a validation step -run_step() { - local step_name=$1 - local step_command=$2 - - echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" - echo -e "${BLUE}▶ ${step_name}${NC}" - echo -e "${BLUE}━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━${NC}" - - if eval "$step_command"; then - echo -e "${GREEN}✅ ${step_name} - PASSED${NC}" - echo "" - return 0 - else - echo -e "${RED}❌ ${step_name} - FAILED${NC}" - echo "" - return 1 - fi -} - -# Track failures -FAILED_STEPS=() - -# Step 1: Code Formatting Check (Ruff) -# - default: check only -# - --fix: write changes -if ! run_step "1. Code Formatting (Ruff)" "if \$FIX_MODE; then ruff format laborious/ tests/; else ruff format --check laborious/ tests/ e2e/; fi"; then - FAILED_STEPS+=("Code Formatting") -fi - -# Step 2: Linting (Ruff) -# - default: check only -# - --fix: apply autofixes -if ! run_step "2. Code Linting (Ruff)" "if \$FIX_MODE; then ruff check --fix laborious/ tests/; else ruff check laborious/ tests/ e2e/; fi"; then - FAILED_STEPS+=("Linting") -fi - -# Step 3: Type Checking (mypy) -if ! run_step "3. Type Checking (mypy)" "mypy laborious/"; then - FAILED_STEPS+=("Type Checking") -fi - -# Step 4: Security Analysis (Bandit) -if ! run_step "4. Security Analysis (Bandit)" "bandit -c pyproject.toml -r laborious/ -ll -q"; then - FAILED_STEPS+=("Security Analysis") -fi - -# Step 5: Unit Tests (pytest) -if ! run_step "5. Unit Tests (pytest)" "pytest tests/ --cov=laborious --cov-report=term-missing --cov-report=xml --cov-report=html --cov-fail-under=80 -q"; then - FAILED_STEPS+=("Unit Tests") -fi - -# Summary -echo -e "${BLUE}╔════════════════════════════════════════════════════════╗${NC}" -echo -e "${BLUE}║ Validation Summary ║${NC}" -echo -e "${BLUE}╚════════════════════════════════════════════════════════╝${NC}" -echo "" - -if [ ${#FAILED_STEPS[@]} -eq 0 ]; then - echo -e "${GREEN}✅ All validation checks passed!${NC}" - echo -e "${GREEN} Your code is ready for commit/deployment.${NC}" - echo "" - exit 0 -else - echo -e "${RED}❌ Validation failed for the following steps:${NC}" - for step in "${FAILED_STEPS[@]}"; do - echo -e "${RED} • ${step}${NC}" - done - echo "" - echo -e "${YELLOW}💡 Tips:${NC}" - echo -e "${YELLOW} • Run 'ruff format laborious/ tests/' to auto-fix formatting${NC}" - echo -e "${YELLOW} • Run 'ruff check --fix laborious/ tests/' to auto-fix linting issues${NC}" - echo -e "${YELLOW} • Review mypy errors and add type hints where needed${NC}" - echo -e "${YELLOW} • Check bandit warnings for security issues${NC}" - echo -e "${YELLOW} • Fix failing tests or improve test coverage${NC}" - echo "" - exit 1 -fi \ No newline at end of file From d43f08d272c4ca2a24f9830725897957384abd84 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 09:22:08 -0300 Subject: [PATCH 03/51] SIENTIAPDE-1712 Enhance README and Implement Drift Detection and Metrics Workflows MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Added new sections in README for Drift Workflow and Simple Metrics Workflow, detailing their execution flows and functionalities. - Introduced `drift.py` for data drift detection, comparing current data against reference datasets. - Added `simple_metrics.py` for calculating regression metrics (RMSE, MSE, MAE, R²). - Updated `values.yaml` to include configuration for MinIO retention hours and offload threshold. - Refactored `minio_dataframe_payload.py` to use the new offload threshold environment variable. - Adjusted tests to reflect changes in environment variable handling for MinIO offload threshold. --- README.md | 217 ++++++++++++++---- .../utils/models/minio_dataframe_payload.py | 2 +- tests/conftest.py | 4 +- values.yaml | 4 + 4 files changed, 176 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index a02f756..c9c37d6 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,8 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc - [Prediction Process Workflow](#2-prediction-process-workflow-prediction_processpy) - [Format and Export Prediction Workflow](#3-format-and-export-prediction-workflow-format_and_export_predictionpy) - [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy) + - [Drift Workflow](#5-drift-workflow-driftpy) + - [Simple Metrics Workflow](#6-simple-metrics-workflow-simple_metricspy) - [Installation & Setup](#installation--setup) - [Prerequisites](#prerequisites) - [Environment Setup](#environment-setup) @@ -77,13 +79,16 @@ A comprehensive, Temporal-based ML orchestration system for industrial data proc ### Advanced Capabilities - **Incremental Data Processing**: Timestamp-based loading to avoid reprocessing - **Configurable Data Retention**: Model retention policies with automatic cleanup +- **MinIO Payload Offload**: Automatic offload of large DataFrames to MinIO with retention cleanup +- **Data Drift Detection**: Univariate and multivariate drift monitoring against reference data +- **Regression Metrics**: Automated RMSE, MSE, MAE, R² calculation and export - **Notification System**: Integrated alerting via MongoDB - **Scalable Architecture**: Kubernetes-ready with horizontal scaling - **Model Retraining**: Automated retraining workflows with production model updates ### Development & Quality Assurance - **Code Quality Tools**: Ruff (lint/format), mypy (types), Bandit (security) -- **Automated Validation**: `validate.sh` and CI quality gates +- **Automated Validation**: CI quality gates and individual tool commands - **Comprehensive Testing**: pytest with async support and high coverage - **Type Safety**: Static type checking with mypy - **Coverage Visualization**: Coverage Gutters integration @@ -129,6 +134,8 @@ Laborious uses a Temporal-based architecture with strong separation of concerns - `sub_workflows/prediction_process.py`: Core prediction pipeline - `sub_workflows/format_and_export_prediction.py`: Formatting and export - `minimal_retrain.py`: Automated model retraining and production update +- `drift.py`: Data drift detection and monitoring +- `simple_metrics.py`: Regression metrics calculation (RMSE, MSE, MAE, R²) #### **Activities (`laborious/activities/`)** - `gates.py`: Data quality validation, filtering, and data formatting operations @@ -139,16 +146,26 @@ Laborious uses a Temporal-based architecture with strong separation of concerns - MLFlow model transformation and prediction - Model retraining and production updates - Reference data retrieval from MLflow Model Registry +- `storage.py`: PostgreSQL queries and MinIO-aware data loading + - `load_query_with_minio_offload`: SQL load with automatic MinIO offload + - `export_payload_to_postgres`: Resolve MinIO payloads and export to Postgres + - `cleanup_minio_objects_expired`: Retention-based MinIO object cleanup + - `query_to_minio`: Legacy parquet upload for retraining data +- `model_metrics.py`: Drift detection and regression metrics + - Univariate and multivariate drift calculation + - Simple metrics (RMSE, MSE, MAE, R²) - `opc.py`: OPC UA export to industrial systems (optional) - `api.py`: PI Web API export operations (optional) - Prediction and confidence data writing to PI Web API - Error handling and notification integration -- `activities.py`: Aggregates activity interfaces +- `activities.py`: Aggregates all activity interfaces (Storage, MLFlow, Gates, OPC, ModelMetrics, API) #### **Data Services (`laborious/utils/`)** - `connectors_config.py`: Env-driven configuration builders +- `models/minio_dataframe_payload.py`: MinIO-offloaded DataFrame payload model - `repository/model_repository.py`: MLFlow operations and retraining - `repository/opc_repository.py`: OPC communication and writes +- `repository/minio_manager.py`: MinIO object storage operations - `filters/conditional_filters.py` and `filters/mlflow_filters.py` ### Data Flow Architecture @@ -167,6 +184,18 @@ Training Data → Model Retraining → Quality Validation → Production Update → Notification & Monitoring ``` +#### **3. Drift Detection Pipeline** +``` +Target Data (PostgreSQL) + Reference Data (MLFlow) → +Drift Calculation (univariate + multivariate) → PostgreSQL Export +``` + +#### **4. Simple Metrics Pipeline** +``` +Predictions + Targets (PostgreSQL JOIN) → +Metrics Calculation (RMSE, MSE, MAE, R²) → PostgreSQL Export +``` + ### Security Architecture #### **Authentication & Authorization** @@ -262,20 +291,21 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M - **Prediction Export**: Delegates prediction formatting and export operations #### Execution Flow -1. **Timestamp Retrieval**: Gets the last processed timestamp for incremental processing -2. **Input Data Gate**: Applies configured filters for data quality validation -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 +1. **Input Data Gate**: Applies configured filters for data quality validation +2. **Path Decision**: Determines processing path based on filter results +3. **MLFlow Transform**: Requests data transformation using MLFlow models +4. **Response Validation**: Filters transform responses for quality assurance +5. **Content Validation**: Filters transformed data content for quality check 6. **MLFlow Prediction**: Executes prediction using transformed data -7. **Content Validation**: Filters prediction responses for final quality check +7. **Prediction Response Validation**: Filters prediction responses for final quality check 8. **Export Delegation**: Delegates to FormatAndExportPrediction workflow +9. **MinIO Cleanup**: Cleans up expired offloaded payloads (if any, in `finally` block) #### Key Features - **Configurable Quality Gates**: Multiple filter types with policy-based configuration - **Flexible Path Handling**: Configurable decision paths (STOP, CONTINUE, REPEAT) - **MLFlow Integration**: Comprehensive model management and inference -- **Incremental Processing**: Timestamp-based data processing optimization +- **MinIO Cleanup**: Automatic retention-based cleanup of offloaded payloads - **Comprehensive Monitoring**: Detailed metrics and error reporting #### Input Parameters @@ -320,12 +350,12 @@ The **PredictionProcess** workflow implements the core prediction pipeline for M #### Architecture Diagram ```mermaid flowchart LR - A[1. get_last_timestamp] --> B[2. input_gate] --> C[3. request_transform] --> D[4. mlflow_response_gate] --> E[5. mlflow_content_gate] --> F[6. request_predict] --> G[7. mlflow_response_gate] --> H[8. format_and_export_prediction🔃] + A[1. input_gate] --> B[2. request_transform] --> C[3. mlflow_response_gate] --> D[4. mlflow_content_gate] --> E[5. request_predict] --> F[6. mlflow_response_gate] --> G[7. format_and_export_prediction🔃] + G --> H[8. cleanup_minio_objects_expired] - A -.-> Redis[(Redis)] - C -.-> MLFlow[MLFlow] - F -.-> MLFlow[MLFlow] - G -.-> Filters[MLFlow Filters] + B -.-> MLFlow[MLFlow] + E -.-> MLFlow[MLFlow] + H -.-> MinIO[(MinIO)] ``` ### 3. Format and Export Prediction Workflow (`format_and_export_prediction.py`) @@ -403,6 +433,76 @@ flowchart LR D -.-> PostgreSQL[(PostgreSQL)] ``` +### 5. Drift Workflow (`drift.py`) + +The **Drift** workflow detects data drift by comparing current data against a reference dataset from the MLflow Model Registry. + +#### Execution Flow +1. **Data Loading**: Loads target data and reference data in parallel +2. **Drift Calculation**: Calculates univariate and multivariate drift metrics +3. **Data Export**: Exports drift metrics to PostgreSQL + +#### Architecture Diagram +```mermaid +flowchart LR + A[1. load_custom_query] --> C[3. calculate_drift] --> D[4. export_data_to_postgres] + B[2. get_reference_data] --> C + + A -.-> Database[(Database)] + B -.-> MLFlow[MLFlow] + D -.-> PostgreSQL[(PostgreSQL)] +``` + +#### Input Parameters +```json +{ + "schedule_name": "hourly_drift", + "model_name": "temperature_model", + "model_id": "temp_001", + "schema": "sientia_data", + "source_table_name": "laborious_data", + "target_table_name": "drift_metrics", + "interval": 60, + "model_config": { "target": "temperature" }, + "drift_metrics": ["kolmogorov_smirnov", "jensen_shannon", "wasserstein"], + "chunk_period": "min" +} +``` + +### 6. Simple Metrics Workflow (`simple_metrics.py`) + +The **SimpleMetrics** workflow calculates regression metrics (RMSE, MSE, MAE, R²) by comparing predictions against actual target values. + +#### Execution Flow +1. **Data Loading**: Loads prediction vs target data via a JOIN query +2. **Metrics Calculation**: Calculates configured regression metrics +3. **Data Export**: Exports metrics to PostgreSQL + +#### Architecture Diagram +```mermaid +flowchart LR + A[1. load_custom_query] --> B[2. calculate_simple_metrics] --> C[3. export_data_to_postgres] + + A -.-> Database[(Database)] + C -.-> PostgreSQL[(PostgreSQL)] +``` + +#### Input Parameters +```json +{ + "schedule_name": "hourly_metrics", + "model_name": "temperature_model", + "model_id": "temp_001", + "schema": "sientia_data", + "predictions_table_name": "predictions", + "data_table_name": "laborious_data", + "target_table_name": "simple_metrics", + "interval_minutes": 60, + "model_config": { "target": "temperature" }, + "metrics": ["rmse", "mse", "mae", "r2"] +} +``` + ## 📋 Prerequisites - Python 3.11+ @@ -527,8 +627,8 @@ pytest pytest --cov=laborious --cov-report=html # Run specific test categories -pytest tests/activities/ -pytest tests/workflow/ +pytest tests/laborious/activities/ +pytest tests/laborious/workflows/ ``` ### Manual Application Execution @@ -569,18 +669,7 @@ pip install -r requirements-dev.txt ### Complete Validation -Option 1 (recommended): -```bash -./validate.sh -``` -The script runs, in order: -1. Format check (Ruff) -2. Linting (Ruff) -3. Type checking (mypy) -4. Security analysis (Bandit) -5. Tests with coverage (pytest) - -Option 2 (individual commands): +Run each validation step individually: ```bash ruff format --check laborious/ tests/ ruff check laborious/ tests/ @@ -606,7 +695,7 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac ### Best Practices -- Run `./validate.sh` before committing +- Run all validation steps before committing - Use `ruff check --watch` for continuous feedback - Add type hints and tests for new code @@ -615,15 +704,34 @@ The workflow at `.github/workflows/quality-gate.yml` executes validations on eac ### Test Structure ``` tests/ -├── activities/ # Activity implementation tests -│ ├── test_gates.py # Data quality gates and formatting tests -│ ├── test_mlflow.py # MLFlow operations and reference data tests -│ └── ... # Other activity tests -├── workflows/ # Workflow orchestration tests -│ └── subworkflows/ # Sub-workflow tests -│ └── test_format_and_export_prediction.py # Export workflow tests -├── utils/ # Utility function tests -└── integration/ # End-to-end workflow tests +├── conftest.py # Global fixtures and env setup +├── laborious/ +│ ├── activities/ # Activity implementation tests +│ │ ├── test_activities.py # Activities aggregator tests +│ │ ├── test_gates.py # Data quality gates and formatting tests +│ │ ├── test_mlflow.py # MLFlow operations and reference data tests +│ │ ├── test_storage.py # Storage and MinIO offload tests +│ │ ├── test_model_metrics.py # Drift and simple metrics tests +│ │ ├── test_opc.py # OPC operations tests +│ │ └── test_api.py # PI Web API operations tests +│ ├── workflows/ # Workflow orchestration tests +│ │ ├── test_predictions_batch.py +│ │ ├── test_minimal_retrain.py +│ │ ├── test_drift.py +│ │ ├── test_simple_metrics.py +│ │ └── subworkflows/ +│ │ ├── test_prediction_process.py +│ │ └── test_format_and_export_prediction.py +│ └── utils/ # Utility function tests +│ ├── test_connectors_config.py +│ ├── models/ +│ │ └── test_minio_dataframe_payload.py +│ ├── filters/ +│ │ ├── test_conditional_filters.py +│ │ └── test_mlflow_filters.py +│ └── repository/ +│ ├── test_model_repository.py +│ └── test_opc_repository.py ``` ### Test Coverage @@ -643,8 +751,8 @@ pip install pytest pytest-cov pytest-asyncio pytest --cov=laborious --cov-report=html # Run specific test modules -pytest tests/activities/test_gates.py -pytest tests/workflow/test_predictions_batch.py +pytest tests/laborious/activities/test_gates.py +pytest tests/laborious/workflows/test_predictions_batch.py ``` ## 📊 Monitoring and Metrics @@ -716,8 +824,13 @@ The Laborious system exposes comprehensive Prometheus metrics for operational vi | `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No | | `HTTP_SDK_METRICS_PORT` | Temporal SDK metrics port | `9091` | No | | `POD_ID` | Kubernetes pod identifier | `None` | No | -| `SIENTIA_MINIO_RETENTION_HOURS` | Retention window for offloaded MinIO objects | `168` | No | -| `SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES` | Offload threshold for DataFrame-derived payloads | `int(1.5 * 1024 * 1024)` | No | +| `MINIO_ENDPOINT_URL` | MinIO endpoint URL | `http://localhost:9000` | Yes | +| `MINIO_ACCESS_KEY` | MinIO access key | `minioadmin` | Yes | +| `MINIO_SECRET_KEY` | MinIO secret key | `minioadmin` | Yes | +| `MINIO_REGION_NAME` | MinIO region name | `us-east-1` | No | +| `MINIO_DEFAULT_BUCKET` | Default MinIO bucket | `laborious` | No | +| `MINIO_RETENTION_HOURS` | Retention window (hours) for offloaded MinIO objects | `24` | No | +| `SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES` | Offload threshold for DataFrame-derived payloads | `int(1.5 * 1024 * 1024)` | No | @@ -728,7 +841,7 @@ Laborious uses MinIO to prevent Temporal workflow history from carrying very lar Whenever a payload exceeds a configurable size threshold, it is stored as a parquet file in MinIO and the workflow history only keeps a lightweight reference. Notes: -- `SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES` supports: +- `SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES` supports: - Integer bytes (e.g. `"1572864"`) - Float MiB (e.g. `"1.5"`), converted to bytes as `MiB * 1024 * 1024` - Fallback behavior uses `1.5 MiB` when the env var is missing or invalid. @@ -751,7 +864,7 @@ After evaluation, the payload is serialized for Temporal as a flat dict: - `data` is omitted / set to `None`. When an activity needs pandas operations, it resolves references using: -- `MinioDataFramePayload.dataframe_from_wire(...)` +- `MinioDataFramePayload.retrieve(minio_repo)` — downloads from MinIO or returns inline data as a DataFrame #### MinIO Object Naming (Retention Parsing) @@ -952,27 +1065,35 @@ This is the configuration created by the Orchestrator in Temporal. ``` laborious/ ├── activities/ # Temporal activity implementations -│ ├── activities.py # Main activities orchestrator +│ ├── activities.py # Main activities aggregator │ ├── gates.py # Data quality gates and filtering │ ├── mlflow.py # MLFlow model operations +│ ├── storage.py # PostgreSQL queries and MinIO offload +│ ├── model_metrics.py # Drift and regression metrics │ ├── opc.py # OPC server operations │ └── api.py # PI Web API operations ├── workflows/ # Temporal workflow definitions │ ├── predictions_batch.py # Main batch prediction workflow │ ├── minimal_retrain.py # Model retraining workflow +│ ├── drift.py # Data drift detection workflow +│ ├── simple_metrics.py # Regression metrics workflow │ └── sub_workflows/ # Sub-workflow implementations │ ├── prediction_process.py # Core prediction workflow │ └── format_and_export_prediction.py # Export workflow ├── worker/ # Worker implementation -│ └── worker.py # Main worker orchestrator +│ ├── worker.py # Main worker orchestrator +│ └── prepare_worker.py # Worker factory with autoscaling config ├── utils/ # Utility functions -│ ├── connectors_config.py # Database configuration +│ ├── connectors_config.py # Environment-driven config builders +│ ├── models/ # Data models +│ │ └── minio_dataframe_payload.py # MinIO-offloaded DataFrame payload │ ├── filters/ # Data quality filters │ │ ├── conditional_filters.py # Conditional data filters │ │ └── mlflow_filters.py # MLFlow response filters │ └── repository/ # Data access layer │ ├── model_repository.py # MLFlow model operations -│ └── opc_repository.py # OPC server operations +│ ├── opc_repository.py # OPC server operations +│ └── minio_manager.py # MinIO object storage operations ├── metrics.py # Prometheus metrics definitions └── __init__.py ``` diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 4a51bb3..7fd9bc8 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -31,7 +31,7 @@ _OBJECT_TIMESTAMP_PATTERN = re.compile( ) OFFLOAD_THRESHOLD_BYTES = int( - float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1.5')) * 1024 * 1024 + float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024 ) # Relative prefix used for storing offloaded training datasets in MinIO. diff --git a/tests/conftest.py b/tests/conftest.py index 0cf3b48..7671dce 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,9 +2,9 @@ import os import sys from unittest.mock import MagicMock -# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES to int at import-time. +# The production code converts SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to int at import-time. # Tests must set it to a valid integer string to avoid import errors. -os.environ.setdefault('SIENTIA_MINIO_OFFLOAD_THRESHOLD_BYTES', '1') +os.environ.setdefault('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1') class DummyMinioDataFramePayload: diff --git a/values.yaml b/values.yaml index 08fa8d3..f1313d4 100644 --- a/values.yaml +++ b/values.yaml @@ -231,6 +231,10 @@ env: value: "sa-east-1" - name: MINIO_DEFAULT_BUCKET value: "sientia" + - name: MINIO_RETENTION_HOURS + value: "24" + - name: SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES + value: "1.5" - name: PI_WEB_API_BASE_URL value: "https://pivision.votorantimcimentos.com/piwebapi" From add3629272b77ef3cbed3d7d98b120e4c9c3fb27 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 14:37:26 -0300 Subject: [PATCH 04/51] SIENTIAPDE-1712 Update liveness and readiness probe initial delays in values.yaml; change GITHUB_BRANCH to feature/SIENTIAPDE-1712; refactor MinioRepository initialization in activities.py --- laborious/activities/activities.py | 2 +- values.yaml | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index 39ee7d7..9a2c4a5 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -75,7 +75,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): metrics_controller = MetricsController(logger=logger) minio_repository = MinioRepository( - endpoint_url=minio_config['endpoint_url'], + endpoint=minio_config['endpoint_url'], access_key=minio_config['access_key'], secret_key=minio_config['secret_key'], bucket=minio_config['default_bucket'], diff --git a/values.yaml b/values.yaml index f1313d4..d952cde 100644 --- a/values.yaml +++ b/values.yaml @@ -71,7 +71,7 @@ livenessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 660 + initialDelaySeconds: 1260 periodSeconds: 15 timeoutSeconds: 5 failureThreshold: 3 @@ -83,7 +83,7 @@ readinessProbe: - -c - | curl -sf http://localhost:9090/metrics | grep -q '^app_up{.*} 1' - initialDelaySeconds: 600 + initialDelaySeconds: 1200 periodSeconds: 10 timeoutSeconds: 3 failureThreshold: 2 @@ -157,7 +157,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" - name: GITHUB_BRANCH - value: "fix/SIENTIAPDE-1478" + value: "feature/SIENTIAPDE-1712" - name: PYTHON_APP value: "laborious.worker.worker" From 18718418fc451cc3319622a60a2366f38f4fed2c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 14:59:47 -0300 Subject: [PATCH 05/51] SIENTIAPDE-1712 Update MinIO configuration in values.yaml and activities.py; add secure option in connectors_config.py --- laborious/activities/activities.py | 1 + laborious/utils/connectors_config.py | 3 ++- values.yaml | 2 +- 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/laborious/activities/activities.py b/laborious/activities/activities.py index 9a2c4a5..9a03f54 100644 --- a/laborious/activities/activities.py +++ b/laborious/activities/activities.py @@ -82,6 +82,7 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API): logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, + secure=minio_config['secure'], ) # Initialize parent classes diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 5e3f186..94dd6a1 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -78,7 +78,7 @@ def build_minio_config() -> dict[str, Any]: 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 """ @@ -89,4 +89,5 @@ def build_minio_config() -> dict[str, Any]: 'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'), 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'), 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), + 'secure': getenv('MINIO_SECURE', 'false') == 'true', } diff --git a/values.yaml b/values.yaml index d952cde..0c4a997 100644 --- a/values.yaml +++ b/values.yaml @@ -222,7 +222,7 @@ env: value: "1" - name: MINIO_ENDPOINT_URL - value: "http://minio.minio.svc.cluster.local:9000" + value: "minio.minio.svc.cluster.local:9000" - name: MINIO_ACCESS_KEY value: "admin" - name: MINIO_SECRET_KEY From 57b0005c94241b4a3346122131c48719b4f81f81 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 15:06:23 -0300 Subject: [PATCH 06/51] SIENTIAPDE-1712 Update MinIO configuration by changing access key in values.yaml and removing region_name references in multiple files to streamline configuration. --- e2e/conftest.py | 1 - laborious/utils/connectors_config.py | 1 - tests/laborious/activities/test_activities.py | 2 -- tests/laborious/utils/test_connectors_config.py | 2 -- values.yaml | 4 +--- 5 files changed, 1 insertion(+), 9 deletions(-) diff --git a/e2e/conftest.py b/e2e/conftest.py index ca3830e..61f86f2 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -410,7 +410,6 @@ async def test_activities( 'endpoint_url': 'http://localhost:9000', 'access_key': 'test', 'secret_key': 'test', - 'region_name': 'us-east-1', 'default_bucket': 'test-bucket', }, opc_config={}, diff --git a/laborious/utils/connectors_config.py b/laborious/utils/connectors_config.py index 94dd6a1..ed3477a 100644 --- a/laborious/utils/connectors_config.py +++ b/laborious/utils/connectors_config.py @@ -86,7 +86,6 @@ def build_minio_config() -> dict[str, Any]: 'endpoint_url': getenv('MINIO_ENDPOINT_URL', 'http://localhost:9000'), 'access_key': getenv('MINIO_ACCESS_KEY', 'minioadmin'), 'secret_key': getenv('MINIO_SECRET_KEY', 'minioadmin'), - 'region_name': getenv('MINIO_REGION_NAME', 'us-east-1'), 'default_bucket': getenv('MINIO_DEFAULT_BUCKET', 'laborious'), 'retention_hours': int(getenv('MINIO_RETENTION_HOURS', '24')), 'secure': getenv('MINIO_SECURE', 'false') == 'true', diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index 4c7b6e6..dae618e 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -43,7 +43,6 @@ def test___init__( 'endpoint_url': 'localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', - 'region_name': 'us-east-1', 'default_bucket': 'test', 'retention_hours': 24, } @@ -187,7 +186,6 @@ async def test_shutdown( 'endpoint_url': 'localhost:9000', 'access_key': 'minio', 'secret_key': 'minio123', - 'region_name': 'us-east-1', 'default_bucket': 'test', 'retention_hours': 24, } diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index e3ba910..4c56576 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -100,7 +100,6 @@ def test_build_minio_config_with_env_vars(): 'endpoint_url': 'http://test-host', 'access_key': 'test-key', 'secret_key': 'test-secret', - 'region_name': 'test-region', 'default_bucket': 'test-bucket', 'retention_hours': 24, } @@ -116,7 +115,6 @@ def test_build_minio_config_with_defaults(): 'endpoint_url': 'http://localhost:9000', 'access_key': 'minioadmin', 'secret_key': 'minioadmin', - 'region_name': 'us-east-1', 'default_bucket': 'laborious', 'retention_hours': 24, } diff --git a/values.yaml b/values.yaml index 0c4a997..215b9bf 100644 --- a/values.yaml +++ b/values.yaml @@ -226,9 +226,7 @@ env: - name: MINIO_ACCESS_KEY value: "admin" - name: MINIO_SECRET_KEY - value: "FvcxOPX55j" - - name: MINIO_REGION_NAME - value: "sa-east-1" + value: "LiArt4eNmJ" - name: MINIO_DEFAULT_BUCKET value: "sientia" - name: MINIO_RETENTION_HOURS From 8789e6693f9e2900557bfe568b1e87a5aab1bba7 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 15:27:55 -0300 Subject: [PATCH 07/51] SIENTIAPDE-1712 Remove `query_to_minio` method from Storage class and update worker activities to eliminate its usage. This change streamlines the codebase by removing unused functionality related to MinIO queries. --- laborious/activities/storage.py | 69 --------------------------------- laborious/worker/worker.py | 3 -- 2 files changed, 72 deletions(-) diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 9fda091..223001b 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -206,75 +206,6 @@ class Storage(Postgres, MinioManager): return report - @activity.defn(name='query_to_minio') - async def query_to_minio(self, input_data: dict[str, Any]) -> dict[str, Any]: - """ - Execute SQL query, write result as Parquet to MinIO, and return object name. - - Args (input_data): - metadata (dict): Workflow metadata - query (str): SQL query - model_name (str): Model name for object naming - object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain) - - Returns: - dict: { success: bool, object_name: str, uri: str } - """ - - if self.minio_repository is None: - raise ValueError('Minio repository not initialized') - - metadata = input_data.get('metadata', {}) - model_name = input_data.get('model_name') or metadata.get('model_name') or 'unknown' - object_prefix = input_data.get('object_prefix', 'datasets/retrain') - - timestamp = now().strftime(DATETIME_FORMAT_FILENAME) - # Keep a stable model-level layout for minimal_retrain: - # training_datasets// - # Sanitize object_prefix to avoid extra subdirectories in the relative key. - safe_prefix = str(object_prefix).strip().strip('/').replace('/', '_') - filename = f'{safe_prefix}_{timestamp}.parquet' - relative_key = f'training_datasets/{model_name}/{filename}' - bucket = getattr(self.minio_repository, 'bucket', 'streamlit-connectors') - uri = f's3://{bucket}/{relative_key}' - - try: - data = await self.load_custom_query(input_data) - if not data: - self.error('query_to_minio failed: No data returned from query', metadata) - return {'success': False, 'message': 'No data returned from query'} - - # Ensure we have a DataFrame - data = pd.DataFrame(data) - - # Convert DataFrame -> parquet bytes, then upload using the new MinIO interface. - parquet_buffer = BytesIO() - data.to_parquet(parquet_buffer, engine='pyarrow', index=True) - file_bytes = parquet_buffer.getvalue() - - upload_result = await self.minio_repository.upload_file( - file_bytes=file_bytes, - relative_key=relative_key, - metadata=metadata, - ) - - object_key_full = upload_result.get('minio_object_name', relative_key) - uri = f's3://{bucket}/{object_key_full}' - return {'success': True, 'object_key': object_key_full, 'uri': uri} - except Exception as e: - trace = traceback.format_exc() - await self.send_notification_async( - metadata=metadata, - notification_id='ERROR_STORING_QUERY_TO_MINIO', - message=f'Error storing query to MinIO: {e}', - block='query_to_minio', - level=NotificationLevel.ERROR, - attachment_content=trace, - ) - - self.error(trace, metadata) - - return {'success': False, 'message': str(e)} def close(self) -> None: """Close Storage resources (MinIO client and Postgres engine).""" diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 8c63b34..5202f2d 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -152,9 +152,7 @@ async def main(): main_workflow=MinimalRetrain, other_workflows=[], activities=[ - activities.load_custom_query, activities.load_query_with_minio_offload, - activities.query_to_minio, activities.retrain_model, activities.update_production_model, activities.format_retrain_report, @@ -200,7 +198,6 @@ async def main(): activities.format_transformed_data, activities.format_prediction, activities.format_default_prediction, - activities.get_last_timestamp, # OPC activities.write_opc_data, # Postgres / MinIO offload From 67942c45e0eb28fd7347cb1f2e4612eea340a9e3 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 15:52:04 -0300 Subject: [PATCH 08/51] SIENTIAPDE-1712 Refactor MinioDataFramePayload usage across activities - Updated instances of MinioDataFramePayload initialization in Gates, MLFlow, and Storage classes to use the new from_dict method for better data reconstruction from dictionaries. - Enhanced the PredictionProcess workflow to utilize the updated payload handling. - Added passthrough fixtures in tests to accommodate the new from_dict method for consistent testing behavior. --- laborious/activities/gates.py | 8 +-- laborious/activities/mlflow.py | 6 +-- laborious/activities/storage.py | 9 ++-- .../utils/models/minio_dataframe_payload.py | 33 +++++++++++- .../sub_workflows/prediction_process.py | 17 +++--- tests/laborious/activities/test_gates.py | 8 +++ tests/laborious/activities/test_mlflow.py | 8 +++ tests/laborious/activities/test_storage.py | 9 ++++ .../models/test_minio_dataframe_payload.py | 52 +++++++++++++++++++ .../subworkflows/test_prediction_process.py | 9 ++++ .../workflows/test_minimal_retrain.py | 9 ++++ .../workflows/test_predictions_batch.py | 18 ++++--- 12 files changed, 155 insertions(+), 31 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index bd40b53..7acafd3 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -157,7 +157,7 @@ class Gates(MinioManager): self.info('Performing input gate...', metadata) filters = input_data['filters'] - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) path_priority = input_data['path_priority'] @@ -236,7 +236,7 @@ class Gates(MinioManager): filters = input_data['filters'] - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) gate_type = input_data['type'] @@ -327,7 +327,7 @@ class Gates(MinioManager): filters = input_data['filters'] - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) gate_type = input_data['type'] @@ -465,7 +465,7 @@ class Gates(MinioManager): self.info('Formatting transformed data...', metadata) - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) data['timestamp'] = data.index diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 82c664f..27d1b92 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -128,7 +128,7 @@ class MLFlow(MinioManager): metadata = input_data['metadata'] self.info('Transforming data...', metadata) - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) model_name = input_data['model_name'] @@ -222,7 +222,7 @@ class MLFlow(MinioManager): metadata = input_data['metadata'] self.info('Predicting data...', metadata) - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) model_name = input_data['model_name'] @@ -311,7 +311,7 @@ class MLFlow(MinioManager): try: # Payload-based retrain input (inline dict or MinIO offloaded). - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) except Exception as e: diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 223001b..58b60c9 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -8,7 +8,6 @@ with workflow.unsafe.imports_passed_through(): # Extend the Temporal Postgres activities for convenient query -> MinIO export import traceback from datetime import timedelta - from io import BytesIO from typing import Any import pandas as pd @@ -18,7 +17,7 @@ with workflow.unsafe.imports_passed_through(): 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 DATETIME_FORMAT_FILENAME, now + from sientia_do.temporal.constants import now from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload @@ -116,7 +115,7 @@ class Storage(Postgres, MinioManager): Export a payload to PostgreSQL. """ metadata = input_data.get('metadata') - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) return await self.export_data_to_postgres( @@ -142,7 +141,8 @@ class Storage(Postgres, MinioManager): raise ValueError('Minio repository not initialized') metadata = input_data.get('metadata', {}) - prefix = input_data['prefix'] + 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 @@ -206,7 +206,6 @@ class Storage(Postgres, MinioManager): return report - def close(self) -> None: """Close Storage resources (MinIO client and Postgres engine).""" Postgres.close(self) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 7fd9bc8..8463be2 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -81,6 +81,38 @@ class MinioDataFramePayload: object_prefix: str | None = None uri: str | None = None + @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) -> int: """ @@ -118,7 +150,6 @@ class MinioDataFramePayload: except ValueError: return None - @staticmethod def cleanup_prefix(self) -> str | None: """ Return True if cleanup is enabled for this payload. diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 07e42e7..3560d5f 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -38,8 +38,6 @@ class PredictionProcess: 8. Export Delegation: Delegates to FormatAndExportPrediction workflow """ - cleanup_prefixes: set[str] = set() - @workflow.run async def run(self, input_data: dict[str, Any]): """ @@ -90,8 +88,6 @@ class PredictionProcess: model_config = input_data.get('model_config', {}) save_transform = input_data.get('save_transform', True) - prefix = data.cleanup_prefix() - try: await self._run_prediction_pipeline( input_data, @@ -103,13 +99,12 @@ class PredictionProcess: save_transform, ) finally: - if self.cleanup_prefixes: - await workflow.execute_activity_method( - Activities.cleanup_minio_objects_expired, - {**metadata, 'prefix': prefix}, - retry_policy=retry_policy, - start_to_close_timeout=timedelta(minutes=5), - ) + await workflow.execute_activity_method( + Activities.cleanup_minio_objects_expired, + {**metadata, 'data': data}, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(minutes=5), + ) async def _run_prediction_pipeline( self, diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index eaee176..0d90ea6 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -7,6 +7,14 @@ from sientia_do.notifications.models import NotificationLevel from laborious.activities.gates import Gates +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.gates.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + def _minio_payload(retrieve_return, status=None): """ Build a MinioDataFramePayload-like test double with async retrieve. diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index baa06e4..e495be8 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -8,6 +8,14 @@ from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_ from laborious.activities.mlflow import MLFlow +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + @patch('laborious.activities.mlflow.MLFlowRepository') @patch('laborious.activities.mlflow.MinioRepository') def test___init__(mock_minio_repository, mock_mlflow_repository): diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index f0272d1..a1fa8e8 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -8,6 +8,15 @@ from sientia_do.temporal.activities.postgres import Postgres from laborious.activities.storage import Storage + +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.activities.storage.MinioDataFramePayload.from_dict', side_effect=lambda x: x + ): + yield + + metadata = { 'metadata': { 'model_id': 'test_model_id', diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py index 9349264..de1a334 100644 --- a/tests/laborious/utils/models/test_minio_dataframe_payload.py +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -214,3 +214,55 @@ async def test_from_dataframe_offloaded(mock_now): assert result.bucket == 'test-bucket' assert result.uri == 's3://test-bucket/full/key.parquet' minio.upload_file.assert_awaited_once() + + +def test_from_dict_inline(): + raw = { + 'last_timestamp': '2024-01-01T00:00:00+00:00', + 'status': None, + 'data': {'col1': {0: 'val1'}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } + payload = MinioDataFramePayload.from_dict(raw) + assert isinstance(payload, MinioDataFramePayload) + assert payload.last_timestamp == '2024-01-01T00:00:00+00:00' + assert payload.data == {'col1': {0: 'val1'}} + assert payload.object_key is None + + +def test_from_dict_offloaded(): + raw = { + 'last_timestamp': '2024-06-15T10:30:45+00:00', + 'status': {'success': True}, + 'data': None, + 'bucket': 'my-bucket', + 'object_key': 'training_datasets/model/model-initial-2024-06-15_10-30-45.parquet', + 'object_prefix': 'training_datasets/model', + 'uri': 's3://my-bucket/training_datasets/model/model-initial-2024-06-15_10-30-45.parquet', + } + payload = MinioDataFramePayload.from_dict(raw) + assert isinstance(payload, MinioDataFramePayload) + assert payload.data is None + assert payload.bucket == 'my-bucket' + assert payload.object_key == raw['object_key'] + assert payload.object_prefix == 'training_datasets/model' + assert payload.uri == raw['uri'] + assert payload.status == {'success': True} + + +def test_from_dict_minimal_keys(): + raw = {'last_timestamp': '2024-01-01'} + payload = MinioDataFramePayload.from_dict(raw) + assert payload.last_timestamp == '2024-01-01' + assert payload.data is None + assert payload.bucket is None + assert payload.object_key is None + + +def test_from_dict_passthrough_existing_instance(): + original = MinioDataFramePayload(last_timestamp='2024-01-01', data={'a': 1}, bucket='b') + result = MinioDataFramePayload.from_dict(original) + assert result is original diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 0da262b..89a2f65 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -6,6 +6,15 @@ from laborious.activities.activities import Activities from laborious.workflows.sub_workflows.prediction_process import PredictionProcess +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.workflows.sub_workflows.prediction_process.MinioDataFramePayload.from_dict', + side_effect=lambda x: x, + ): + yield + + @fixture def prediction_process(): return PredictionProcess() diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index be49c89..9af1331 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -6,6 +6,15 @@ from laborious.activities.activities import Activities from laborious.workflows.minimal_retrain import MinimalRetrain +@fixture(autouse=True) +def _passthrough_from_dict(): + with patch( + 'laborious.workflows.minimal_retrain.MinioDataFramePayload.from_dict', + side_effect=lambda x: x, + ): + yield + + @fixture def minimal_retrain() -> MinimalRetrain: return MinimalRetrain() diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 8ebd040..7164b7a 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, call, patch from pytest import fixture, mark @@ -22,12 +22,15 @@ metadata = { @mark.asyncio +@patch( + 'laborious.workflows.predictions_batch.MinioDataFramePayload.from_dict', side_effect=lambda x: x +) @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) -async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): - workflow_mock.execute_activity_method.return_value = { - 'success': True, - 'data': {'col': ['test_data']}, - } +async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch: PredictionsBatch): + activity_return = MagicMock() + activity_return.cleanup_prefix.return_value = None + workflow_mock.execute_activity_method.return_value = activity_return + input_data = { 'schedule_name': 'test_schedule', 'model_name': 'test_model', @@ -62,7 +65,8 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch ) prediction_input = { 'metadata': metadata, - 'data': {'success': True, 'data': {'col': ['test_data']}}, + 'data': activity_return, + 'cleanup_prefix': activity_return.cleanup_prefix(), 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'transform_table_name': input_data['transform_table_name'], From 52bca3f09dea092d6d83c6a72bd555a9a285ea8c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 15:53:54 -0300 Subject: [PATCH 09/51] SIENTIAPDE-1712 Fix type hint in from_dict method of MinioDataFramePayload class for improved clarity and compatibility with type checking. --- laborious/utils/models/minio_dataframe_payload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 8463be2..9d2fe48 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -82,7 +82,7 @@ class MinioDataFramePayload: uri: str | None = None @classmethod - def from_dict(cls, raw: dict[str, Any] | 'MinioDataFramePayload') -> 'MinioDataFramePayload': + def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload': """ Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization. From e35e27bfb2320049bf1ecbf5eced72b67d814994 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 16:00:11 -0300 Subject: [PATCH 10/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Update type hints in PredictionProcess to improve clarity and enforce data structure consistency. Changed `data` parameter to a dictionary type and updated the way `last_timestamp` is accessed. --- laborious/workflows/sub_workflows/prediction_process.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 3560d5f..2c20ce8 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -110,13 +110,13 @@ class PredictionProcess: self, input_data: dict[str, Any], metadata: dict[str, Any], - data: Any, - model_id: 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 + last_timestamp = data['last_timestamp'] # Apply input data quality gates gate_input = { From 2979f2dd5a7399c3797d07e6f0c49c60d8d2d9c5 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 20 Mar 2026 16:59:04 -0300 Subject: [PATCH 11/51] SIENTIAPDE-1712 Update environment variables in values.yaml and enhance metric labels in metrics.py - Changed POSTGRES_USER and POSTGRES_PASSWORD values in values.yaml for improved security. - Added 'runtime' label to metrics in metrics.py for better environment identification. - Updated CORE_LABELS to include 'runtime' for consistency across metrics. - Modified type hint for data parameter in PredictionProcess to use a dictionary for better clarity. - Adjusted tests to reflect changes in core labels and MinIO configuration. --- laborious/metrics.py | 5 +- .../sub_workflows/prediction_process.py | 2 +- tests/laborious/activities/test_activities.py | 5 +- tests/laborious/activities/test_api.py | 17 ++++ tests/laborious/activities/test_gates.py | 41 +++----- tests/laborious/activities/test_storage.py | 95 ++++--------------- .../laborious/utils/test_connectors_config.py | 2 + .../subworkflows/test_prediction_process.py | 14 +-- .../workflows/test_minimal_retrain.py | 9 -- .../workflows/test_predictions_batch.py | 7 +- values.yaml | 4 +- 11 files changed, 67 insertions(+), 134 deletions(-) diff --git a/laborious/metrics.py b/laborious/metrics.py index cdd44db..6b3c30a 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -18,6 +18,7 @@ Key Metric Categories: 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 @@ -35,8 +36,8 @@ APP_UP = Gauge( ['pod_id'], ) -# Core labels used across multiple metrics -CORE_LABELS = ['pod_id', 'model_name', 'workflow_name'] +# Core labels used across multiple laborious metrics (aligned with ``SientiaMonitoring.labels`` subset) +CORE_LABELS = ['pod_id', 'runtime', 'model_name', 'workflow_name'] # Prediction operation metrics PREDICTIONS_WRITTEN_COUNT = Counter( diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 2c20ce8..00ed167 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -244,7 +244,7 @@ class PredictionProcess: async def path_flag_handler( self, - data: MinioDataFramePayload, + data: dict[str, Any], path_flag: str, input_data: dict, confidence: int, diff --git a/tests/laborious/activities/test_activities.py b/tests/laborious/activities/test_activities.py index dae618e..8e6cd78 100644 --- a/tests/laborious/activities/test_activities.py +++ b/tests/laborious/activities/test_activities.py @@ -45,6 +45,7 @@ def test___init__( 'secret_key': 'minio123', 'default_bucket': 'test', 'retention_hours': 24, + 'secure': False, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} @@ -144,13 +145,14 @@ def test___init__( ) mock_minio_repository.assert_called_once_with( - endpoint_url=minio_config['endpoint_url'], + 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=mock_metrics_controller.return_value, + secure=minio_config['secure'], ) @@ -188,6 +190,7 @@ async def test_shutdown( 'secret_key': 'minio123', 'default_bucket': 'test', 'retention_hours': 24, + 'secure': False, } mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'} diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py index e1e21a9..554459e 100644 --- a/tests/laborious/activities/test_api.py +++ b/tests/laborious/activities/test_api.py @@ -78,7 +78,9 @@ def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_clien 'get_core_labels', return_value={ 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', + 'workflow_name': 'test_workflow', 'operation_type': '-', }, ): @@ -86,6 +88,12 @@ def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_clien metadata=metadata['metadata'], operation_type=None ) assert 'operation_type' not in labels + assert labels == { + 'pod_id': 'test_pod', + 'runtime': 'local', + 'model_name': 'test_model', + 'workflow_name': 'test_workflow', + } @patch('laborious.activities.api.PIWebAPIClient') @@ -105,7 +113,9 @@ def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client): 'get_core_labels', return_value={ 'pod_id': 'test_pod', + 'runtime': 'k8s', 'model_name': 'test_model', + 'workflow_name': 'test_workflow', 'operation_type': 'write', }, ): @@ -113,6 +123,7 @@ def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client): metadata=metadata['metadata'], operation_type='write' ) assert labels['operation_type'] == 'write' + assert labels['runtime'] == 'k8s' def test__init__(): @@ -152,6 +163,7 @@ def api(mock_pi_web_api_client): api_instance.get_core_labels = MagicMock( return_value={ 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -334,6 +346,7 @@ async def test_process_pi_web_api_response_success(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -367,6 +380,7 @@ async def test_process_pi_web_api_response_with_errors(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -395,6 +409,7 @@ async def test_process_pi_web_api_response_missing_tags(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -427,6 +442,7 @@ async def test_process_pi_web_api_response_missing_webid(api): tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } @@ -455,6 +471,7 @@ async def test_process_pi_web_api_response_missing_tag_name(api): tags = {'tag1': 'web_id_1'} core_labels = { 'pod_id': 'test_pod', + 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', } diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 0d90ea6..367f5b4 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -820,15 +820,17 @@ async def test_write_metrics(mock_metrics, gates_activity): 'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}}, } await gates_activity.write_metrics(input_data) + core_tags = { + 'pod_id': gates_activity.pod_id, + 'runtime': gates_activity.runtime, + 'model_name': metadata['metadata']['model_name'], + 'workflow_name': metadata['metadata']['workflow_name'], + } gates_activity.emit_metric.assert_has_calls( [ call( metric_object=mock_metrics.PREDICTIONS_WRITTEN_COUNT, - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, ), ] ) @@ -837,11 +839,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_CONFIDENCE_MONITOR, method='set', - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, value=0.9, ), ] @@ -851,11 +849,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR, method='observe', - tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], - }, + tags=core_tags, value=0.1, ), ] @@ -865,9 +859,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag1', }, @@ -880,9 +872,7 @@ async def test_write_metrics(mock_metrics, gates_activity): metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag1', }, @@ -895,9 +885,7 @@ async def test_write_metrics(mock_metrics, gates_activity): call( metric_object=mock_metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag2', }, @@ -910,9 +898,7 @@ async def test_write_metrics(mock_metrics, gates_activity): metric_object=mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': gates_activity.pod_id, - 'model_name': metadata['metadata']['model_name'], - 'workflow_name': metadata['metadata']['workflow_name'], + **core_tags, 'opc_server_id': 'server1', 'tag': 'tag2', }, @@ -943,6 +929,7 @@ async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_act method='observe', tags={ 'pod_id': gates_activity.pod_id, + 'runtime': gates_activity.runtime, 'model_name': metadata['metadata']['model_name'], 'workflow_name': metadata['metadata']['workflow_name'], 'opc_server_id': 'server1', diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index a1fa8e8..9e76c5f 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -117,80 +117,6 @@ def test___init___done_repository(mock_minio_repository, storage): assert storage.minio_repository is not None -@mark.asyncio -async def test_query_to_minio_minio_repository_not_initialized(storage): - storage.minio_repository = None - - with raises(ValueError) as e: - await storage.query_to_minio({}) - - assert str(e.value) == 'Minio repository not initialized' - - -@mark.asyncio -async def test_query_to_minio_not_data(storage): - storage.load_custom_query = AsyncMock(return_value=None) - result = await storage.query_to_minio({}) - - storage.load_custom_query.assert_called_once_with({}) - assert result['success'] is False - assert result['message'] == 'No data returned from query' - - -@mark.asyncio -@patch('laborious.activities.storage.pd.DataFrame') -@patch('laborious.activities.storage.now') -async def test_query_to_minio_success(now, dataframe, storage): - data = [{'a': 1}, {'a': 2}, {'a': 3}] - storage.load_custom_query = AsyncMock(return_value=data) - now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0) - storage.minio_repository.upload_file = AsyncMock( - return_value={ - 'minio_object_name': 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' - } - ) - storage.minio_repository.bucket = 'test' - - result = await storage.query_to_minio({'object_prefix': 'test', **metadata}) - - dataframe.assert_called_once_with(data) - - storage.minio_repository.upload_file.assert_called_once_with( - file_bytes=ANY, - relative_key='training_datasets/test_model/test_2024-01-01_00-00-00.parquet', - metadata=metadata['metadata'], - ) - - assert result['success'] is True - assert ( - result['object_key'] - == 'sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' - ) - assert ( - result['uri'] - == 's3://test/sientia/streamlit-connectors/training_datasets/test_model/test_2024-01-01_00-00-00.parquet' - ) - - -@mark.asyncio -async def test_query_to_minio_error(storage): - storage.send_notification = MagicMock() - storage.send_notification_async = AsyncMock() - - storage.load_custom_query = AsyncMock(side_effect=Exception('test')) - result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'}) - assert result['success'] is False - assert result['message'] == 'test' - storage.send_notification_async.assert_called_once_with( - metadata=metadata['metadata'], - notification_id='ERROR_STORING_QUERY_TO_MINIO', - message='Error storing query to MinIO: test', - block='query_to_minio', - level=NotificationLevel.ERROR, - attachment_content=ANY, - ) - - def test_close(storage): storage.minio_repository = MagicMock() @@ -275,8 +201,11 @@ async def test_cleanup_minio_objects_expired(mock_now, storage): storage.minio_repository.delete_file = AsyncMock() storage.send_notification_async = AsyncMock() + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' + result = await storage.cleanup_minio_objects_expired( - {**metadata, 'prefix': 'training_datasets/m'} + {**metadata, 'data': data_mock} ) assert result['deleted_count'] == 1 @@ -326,8 +255,10 @@ async def test_export_payload_to_postgres(storage): async def test_cleanup_minio_objects_expired_minio_not_initialized(storage): storage.minio_repository = None + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'test' with raises(ValueError, match='Minio repository not initialized'): - await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'}) + await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) @mark.asyncio @@ -340,7 +271,9 @@ async def test_cleanup_minio_objects_expired_unparseable_key(mock_now, storage): storage.minio_repository.delete_file = AsyncMock() storage.send_notification_async = AsyncMock() - result = await storage.cleanup_minio_objects_expired({**metadata, 'prefix': 'test'}) + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'test' + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) assert result['deleted_count'] == 0 assert result['failed_count'] == 0 @@ -356,8 +289,10 @@ async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage): storage.minio_repository.delete_file = AsyncMock(side_effect=Exception('delete error')) storage.send_notification_async = AsyncMock() + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' result = await storage.cleanup_minio_objects_expired( - {**metadata, 'prefix': 'training_datasets/m'} + {**metadata, 'data': data_mock} ) assert result['deleted_count'] == 0 @@ -375,8 +310,10 @@ async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storag storage.send_notification_async = AsyncMock() storage.error = MagicMock() + data_mock = MagicMock() + data_mock.cleanup_prefix.return_value = 'training_datasets/m' result = await storage.cleanup_minio_objects_expired( - {**metadata, 'prefix': 'training_datasets/m'} + {**metadata, 'data': data_mock} ) assert result['deleted_count'] == 0 diff --git a/tests/laborious/utils/test_connectors_config.py b/tests/laborious/utils/test_connectors_config.py index 4c56576..b94bb8d 100644 --- a/tests/laborious/utils/test_connectors_config.py +++ b/tests/laborious/utils/test_connectors_config.py @@ -102,6 +102,7 @@ def test_build_minio_config_with_env_vars(): 'secret_key': 'test-secret', 'default_bucket': 'test-bucket', 'retention_hours': 24, + 'secure': False, } @@ -117,4 +118,5 @@ def test_build_minio_config_with_defaults(): 'secret_key': 'minioadmin', 'default_bucket': 'laborious', 'retention_hours': 24, + 'secure': False, } diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 89a2f65..db75b8c 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -38,7 +38,7 @@ async def test_run(workflow_mock, prediction_process): # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -199,7 +199,7 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -251,7 +251,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -336,7 +336,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -441,7 +441,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -784,7 +784,7 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.last_timestamp = '2024-01-01' + data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() input_data = { 'metadata': metadata, 'data': data_payload, @@ -816,7 +816,7 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): workflow_mock.execute_activity_method.assert_any_call( Activities.cleanup_minio_objects_expired, - {**metadata, 'prefix': 'training_datasets/test'}, + {**metadata, 'data': data_payload}, retry_policy=ANY, start_to_close_timeout=ANY, ) diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index 9af1331..be49c89 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -6,15 +6,6 @@ from laborious.activities.activities import Activities from laborious.workflows.minimal_retrain import MinimalRetrain -@fixture(autouse=True) -def _passthrough_from_dict(): - with patch( - 'laborious.workflows.minimal_retrain.MinioDataFramePayload.from_dict', - side_effect=lambda x: x, - ): - yield - - @fixture def minimal_retrain() -> MinimalRetrain: return MinimalRetrain() diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 7164b7a..279499f 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -22,13 +22,9 @@ metadata = { @mark.asyncio -@patch( - 'laborious.workflows.predictions_batch.MinioDataFramePayload.from_dict', side_effect=lambda x: x -) @patch('laborious.workflows.predictions_batch.workflow', new_callable=AsyncMock) -async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch: PredictionsBatch): +async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch): activity_return = MagicMock() - activity_return.cleanup_prefix.return_value = None workflow_mock.execute_activity_method.return_value = activity_return input_data = { @@ -66,7 +62,6 @@ async def test_run(workflow_mock: AsyncMock, mock_from_dict, predictions_batch: prediction_input = { 'metadata': metadata, 'data': activity_return, - 'cleanup_prefix': activity_return.cleanup_prefix(), 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'transform_table_name': input_data['transform_table_name'], diff --git a/values.yaml b/values.yaml index 215b9bf..e7c32f9 100644 --- a/values.yaml +++ b/values.yaml @@ -167,9 +167,9 @@ env: - name: POSTGRES_PORT value: "5432" - name: POSTGRES_USER - value: "sientia" + value: "postgres" - name: POSTGRES_PASSWORD - value: "sientia" + value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3" - name: POSTGRES_DBNAME value: "sientia" - name: POSTGRES_MIN_CONNECTIONS From e1d1e3d43d1404e5de22fde22d77a4673d9cf263 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 08:06:01 -0300 Subject: [PATCH 12/51] SIENTIAPDE-1712 Refactor imports in gates.py and mlflow.py for improved organization - Removed unnecessary import statement in mlflow.py and re-added it in a more appropriate location. - Cleaned up the workflow metadata assignment in gates.py for better readability. --- laborious/activities/gates.py | 2 +- laborious/activities/mlflow.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 7acafd3..0b747f9 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -479,7 +479,7 @@ class Gates(MinioManager): minio_repo=self.minio_repository, model_name=input_data['model_name'], operation='transform', - workflow_metadata=metadata, + workflow_metadata=metadata ) @activity.defn(name='format_prediction') diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 27d1b92..e6c75e1 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,6 +1,6 @@ from temporalio import activity, workflow -from laborious.utils.repository.minio_manager import MinioManager + with workflow.unsafe.imports_passed_through(): import traceback @@ -20,7 +20,7 @@ with workflow.unsafe.imports_passed_through(): now, ) from sientia_do.utils.formatters import create_sample_dict - + from laborious.utils.repository.minio_manager import MinioManager from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload from laborious.utils.repository.model_repository import MLFlowRepository From 2fa75060754681502c1f9fec1d36c4e19966d09c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 08:35:28 -0300 Subject: [PATCH 13/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Add timestamp column to processed data in MLFlow and improve dataframe validation in MinioDataFramePayload - Added a 'timestamp' column to the processed data in the MLFlow class for better tracking of data entries. - Updated the validation check in MinioDataFramePayload to handle None values for the dataframe more explicitly. --- laborious/activities/mlflow.py | 2 ++ laborious/utils/models/minio_dataframe_payload.py | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index e6c75e1..e6522f0 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -149,6 +149,8 @@ class MLFlow(MinioManager): data.columns.name = None data.index.name = None + data['timestamp'] = data.index + self.debug(f'Processed input data: \n {data.to_csv()}', metadata) # Request transformation from MLFlow model diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 9d2fe48..c49218d 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -194,7 +194,7 @@ class MinioDataFramePayload: MinioDataFramePayload: Instance with data and/or MinIO fields set. """ - if not dataframe or dataframe.empty: + if dataframe is None or dataframe.empty: return cls( data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status ) From 62b885afaeb6a219bc8a671f5a96b00646c95ff2 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 09:26:49 -0300 Subject: [PATCH 14/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Add last_timestamp parameter to MLFlow and Gates activities for enhanced tracking - Introduced last_timestamp parameter in the MLFlow and Gates classes to improve tracking of data processing times. - Updated MinioDataFramePayload to handle last_timestamp, ensuring it defaults to the maximum timestamp from the dataframe if not provided. --- laborious/activities/gates.py | 3 ++- laborious/activities/mlflow.py | 4 ++++ laborious/utils/models/minio_dataframe_payload.py | 5 ++++- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 0b747f9..f9eeb16 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -479,7 +479,8 @@ class Gates(MinioManager): minio_repo=self.minio_repository, model_name=input_data['model_name'], operation='transform', - workflow_metadata=metadata + workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, ) @activity.defn(name='format_prediction') diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index e6522f0..ea97f3b 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -178,6 +178,7 @@ class MLFlow(MinioManager): operation='transform', status=response_data, workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, ) return await MinioDataFramePayload.from_dataframe( @@ -189,6 +190,7 @@ class MLFlow(MinioManager): status={ 'success': True, }, + last_timestamp=payload.last_timestamp, ) @activity.defn(name='request_predict') @@ -260,6 +262,7 @@ class MLFlow(MinioManager): operation='predict', status=response_data, workflow_metadata=metadata, + last_timestamp=payload.last_timestamp, ) return await MinioDataFramePayload.from_dataframe( @@ -271,6 +274,7 @@ class MLFlow(MinioManager): status={ 'success': True, }, + last_timestamp=payload.last_timestamp, ) @activity.defn(name='retrain_model') diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index c49218d..7e0d865 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -173,6 +173,7 @@ class MinioDataFramePayload: operation: OperationKind, status: dict[str, Any] | None = None, workflow_metadata: dict | None = None, + last_timestamp: str | None = None, ) -> 'MinioDataFramePayload': """ Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO. @@ -199,7 +200,9 @@ class MinioDataFramePayload: data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status ) - last_timestamp = max(dataframe['timestamp'].values.tolist()) + + if last_timestamp is None: + last_timestamp = max(dataframe['timestamp'].values.tolist()) if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES: return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp) From 3c45cbbc8a49c3d3c5e55dcdb1b4f8af6b625551 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 09:42:24 -0300 Subject: [PATCH 15/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Add status parameter to MinioDataFramePayload for enhanced data tracking --- laborious/utils/models/minio_dataframe_payload.py | 1 + 1 file changed, 1 insertion(+) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 7e0d865..412e533 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -233,6 +233,7 @@ class MinioDataFramePayload: object_prefix=object_prefix, uri=uri, last_timestamp=last_timestamp, + status=status, ) async def retrieve( From 44558b441561a11383898d593f311e3ee5d36b01 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 09:53:34 -0300 Subject: [PATCH 16/51] SIENTIAPDE-1712 Add status parameter to MinioDataFramePayload for improved data tracking --- laborious/utils/models/minio_dataframe_payload.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 412e533..89d344c 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -205,7 +205,7 @@ class MinioDataFramePayload: last_timestamp = max(dataframe['timestamp'].values.tolist()) if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES: - return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp) + 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) From e17824eb853a72101303660564769a9015b4fd40 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 10:17:57 -0300 Subject: [PATCH 17/51] SIENTIAPDE-1712 Refactor metrics and API handling for improved consistency and clarity - Removed the SIENTIA_CORE_LABELS constant and replaced it with CORE_LABELS for uniformity across metrics. - Updated the API class to ensure operation_type is always included in core labels for PI Web API metrics. - Simplified metric tag handling in the Gates class by consolidating common tags into a single core_tags dictionary. - Enhanced the PredictionProcess class to improve error handling and variable naming for clarity. --- laborious/activities/api.py | 28 +++------ laborious/activities/gates.py | 34 ++++------ laborious/metrics.py | 62 +++---------------- .../sub_workflows/prediction_process.py | 22 ++++--- tests/laborious/activities/test_api.py | 9 ++- 5 files changed, 51 insertions(+), 104 deletions(-) diff --git a/laborious/activities/api.py b/laborious/activities/api.py index 366ecd2..c568e98 100644 --- a/laborious/activities/api.py +++ b/laborious/activities/api.py @@ -74,34 +74,26 @@ class API(SientiaMonitoring): def get_pi_web_api_core_labels( self, metadata: dict[str, Any], - operation_type: str | None = None, + operation_type: str = 'write_pi_web_api_data', ) -> dict[str, Any]: """ - Generate core labels for metrics, optionally including operation_type. + Generate core labels for PI Web API metrics. - This override keeps compatibility with the base implementation while adding - a convenience overload behavior: - - When operation_type is provided, it behaves exactly like the base class, - returning labels that include the operation_type key. - - When operation_type is omitted (None), it removes the operation_type key - from the resulting labels. This is useful for metrics, such as the PI Web - API metrics, that are defined without the operation_type label. + 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 | None): Optional operation type label. If None, the - operation_type key will be removed from the returned labels. + - 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, with operation_type only when provided + dict[str, Any]: Core labels dictionary including operation_type. """ - base_labels = super().get_core_labels( + return super().get_core_labels( metadata=metadata, - operation_type=operation_type or '-', + operation_type=operation_type, ) - if operation_type is None: - base_labels.pop('operation_type', None) - return base_labels def close(self) -> None: """ diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index f9eeb16..75d4b93 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -703,34 +703,30 @@ class Gates(MinioManager): 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={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, ) await self.emit_metric( metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR, method='set', - tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, value=prediction_confidence, ) await self.emit_metric( metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR, method='observe', - tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], - }, + tags=core_tags, value=response_time, ) @@ -741,9 +737,7 @@ class Gates(MinioManager): metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR, method='observe', tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], + **core_tags, 'opc_server_id': server_id, 'tag': tag, }, @@ -753,9 +747,7 @@ class Gates(MinioManager): await self.emit_metric( metric_object=metrics.PREDICTION_OPC_WRITING_COUNT, tags={ - 'pod_id': self.pod_id, - 'model_name': metadata['model_name'], - 'workflow_name': metadata['workflow_name'], + **core_tags, 'opc_server_id': server_id, 'tag': tag, }, diff --git a/laborious/metrics.py b/laborious/metrics.py index 6b3c30a..a67e574 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -26,7 +26,7 @@ Metric Labels: from prometheus_client import Counter, Gauge, Histogram from sientia_do.observability.metrics import ( - CORE_LABELS as SIENTIA_CORE_LABELS, + CORE_LABELS ) # Application health metric @@ -36,9 +36,6 @@ APP_UP = Gauge( ['pod_id'], ) -# Core labels used across multiple laborious metrics (aligned with ``SientiaMonitoring.labels`` subset) -CORE_LABELS = ['pod_id', 'runtime', 'model_name', 'workflow_name'] - # Prediction operation metrics PREDICTIONS_WRITTEN_COUNT = Counter( 'laborious_predictions_written_count', @@ -61,45 +58,6 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram( buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], ) -# ================== MinIO metrics ================== - -MINIO_READ_LAG = Histogram( - 'laborious_minio_read_lag', - 'Lag between the last write to MinIO and the last read from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], - buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], -) - -MINIO_WRITE_LAG = Histogram( - 'laborious_minio_write_lag', - 'Lag between the last write to MinIO and the last read from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], - buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0], -) - -MINIO_READ_COUNT = Counter( - 'laborious_minio_read_count', - 'Number of reads from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_WRITE_COUNT = Counter( - 'laborious_minio_write_count', - 'Number of writes to MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_READ_ERROR_COUNT = Counter( - 'laborious_minio_read_error_count', - 'Number of errors reading from MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) - -MINIO_WRITE_ERROR_COUNT = Counter( - 'laborious_minio_write_error_count', - 'Number of errors writing to MinIO', - [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'], -) # ================== OPC metrics ================== @@ -138,58 +96,58 @@ OPC_CONNECTION_STATUS = Gauge( MODEL_READ_LAG = Histogram( 'laborious_model_read_lag', 'Lag between the start and read of read operations', - SIENTIA_CORE_LABELS, + 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', - SIENTIA_CORE_LABELS, + 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', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_WRITE_COUNT = Counter( 'laborious_model_write_count', 'Number of writes to the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_READ_ERROR_COUNT = Counter( 'laborious_model_read_error_count', 'Number of errors reading from the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_WRITE_ERROR_COUNT = Counter( 'laborious_model_write_error_count', 'Number of errors writing to the model', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_ANALYZE_LAG = Histogram( 'laborious_model_analyze_lag', 'Lag between the start and end of analyze operations', - SIENTIA_CORE_LABELS, + 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', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) MODEL_ANALYZE_ERROR_COUNT = Counter( 'laborious_model_analyze_error_count', 'Number of errors during analyze operations', - SIENTIA_CORE_LABELS, + CORE_LABELS, ) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 00ed167..fbac978 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -98,13 +98,21 @@ class PredictionProcess: model_config, save_transform, ) - finally: 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, @@ -140,7 +148,7 @@ class PredictionProcess: return # Request MLFlow model transformation - response_data = await workflow.execute_local_activity_method( + transformed_data = await workflow.execute_local_activity_method( Activities.request_transform, {**metadata, 'data': data, 'model_name': model_name, 'model_config': model_config}, retry_policy=retry_policy, @@ -153,7 +161,7 @@ class PredictionProcess: { **metadata, 'filters': input_data['mlflow_transform_filters'], - 'data': response_data, + 'data': transformed_data, 'type': 'transform', 'path_priority': input_data['path_priority'], }, @@ -167,8 +175,6 @@ class PredictionProcess: ): return - transformed_data = response_data['content'] - path_flag, confidence, comment = await workflow.execute_local_activity_method( Activities.mlflow_content_gate, { @@ -187,7 +193,7 @@ class PredictionProcess: ): return - response_data = await workflow.execute_local_activity_method( + predicted_data = await workflow.execute_local_activity_method( Activities.request_predict, { **metadata, @@ -205,7 +211,7 @@ class PredictionProcess: { **metadata, 'filters': input_data['mlflow_predict_filters'], - 'data': response_data, + 'data': predicted_data, 'type': 'predict', 'path_priority': input_data['path_priority'], }, @@ -225,7 +231,7 @@ class PredictionProcess: { 'metadata': metadata, 'path_flag': path_flag, - 'data': response_data['content'], + 'data': predicted_data, 'transformed_data': transformed_data if save_transform else None, 'prediction_confidence': confidence, 'timestamp': last_timestamp, diff --git a/tests/laborious/activities/test_api.py b/tests/laborious/activities/test_api.py index 554459e..bf5b9f1 100644 --- a/tests/laborious/activities/test_api.py +++ b/tests/laborious/activities/test_api.py @@ -81,18 +81,17 @@ def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_clien 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', - 'operation_type': '-', + 'operation_type': 'write_pi_web_api_data', }, ): - labels = api_instance.get_pi_web_api_core_labels( - metadata=metadata['metadata'], operation_type=None - ) - assert 'operation_type' not in labels + labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata']) + assert labels['operation_type'] == 'write_pi_web_api_data' assert labels == { 'pod_id': 'test_pod', 'runtime': 'local', 'model_name': 'test_model', 'workflow_name': 'test_workflow', + 'operation_type': 'write_pi_web_api_data', } From ccbed58234adf2b9e0474cbbd982929f0d4bf59e Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 10:24:58 -0300 Subject: [PATCH 18/51] SIENTIAPDE-1712 Update MinioDataFramePayload to reflect changes from training to prediction datasets - Renamed TRAINING_DATASETS_PREFIX to PREDICTION_DATASETS_PREFIX for clarity. - Updated object key naming convention to use prediction datasets directory. --- laborious/utils/models/minio_dataframe_payload.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 89d344c..42b60f4 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -34,9 +34,9 @@ OFFLOAD_THRESHOLD_BYTES = int( float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024 ) -# Relative prefix used for storing offloaded training datasets in MinIO. +# Relative prefix used for storing offloaded prediction datasets in MinIO. # It is also the root directory for retention cleanup listing. -TRAINING_DATASETS_PREFIX = 'training_datasets' +PREDICTION_DATASETS_PREFIX = 'prediction_datasets' OperationKind = Literal['initial', 'transform', 'predict'] @@ -56,11 +56,11 @@ def _build_object_key( tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root). """ # Naming convention: - # - Directory is always `training_datasets/` + # - 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'{TRAINING_DATASETS_PREFIX}/{model_dir}' + prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}' return f'{prefix}/{basename}', prefix From 69cbd08913fc3be206613d34ea45c24188cd0f08 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 10:47:57 -0300 Subject: [PATCH 19/51] SIENTIAPDE-1712 Update payload initialization in Gates class to use from_dict method for improved clarity and consistency --- laborious/activities/gates.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 75d4b93..1bf37f2 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -512,7 +512,7 @@ class Gates(MinioManager): prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) - payload: MinioDataFramePayload = input_data['data'] + payload = MinioDataFramePayload.from_dict(input_data['data']) data = await payload.retrieve(self.minio_repository, metadata) # Create timestamp column from index and reset index From 50a7a73a8d60321220dac35bef38712fa5b7743f Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 12:06:50 -0300 Subject: [PATCH 20/51] SIENTIAPDE-1712 --- laborious/activities/gates.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 1bf37f2..cac6c38 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -233,10 +233,13 @@ class Gates(MinioManager): """ 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(input_data['data']) + payload = MinioDataFramePayload.from_dict(raw_data) data = await payload.retrieve(self.minio_repository, metadata) gate_type = input_data['type'] @@ -244,9 +247,6 @@ class Gates(MinioManager): filter_output = [] - self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}', metadata) - self.debug(f'Filters: {filters}', metadata) - comments = [] status = payload.status or {} From f22cc49b93e84255e1fb84ba4ea54d2be94e8dfd Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 12:28:18 -0300 Subject: [PATCH 21/51] SIENTIAPDE-1712 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 4c30e62..dd4fa1a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.2 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.3 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.7 prometheus-client botocore From 503d9aa485315e23e442f573f815cf15a3bd9d4c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 23 Mar 2026 14:45:21 -0300 Subject: [PATCH 22/51] SIENTIAPDE-1712 Update dependencies and refactor input filter handling for consistency - Updated sientia-dataops-library dependency version from 1.10.3 to 1.10.4 in requirements.txt. - Refactored input filter handling in the Gates class to read policy and config keys in a case-insensitive manner. - Updated test cases to ensure consistency in filter key naming conventions across various scenarios. --- e2e/conftest.py | 32 ++++--- e2e/test_predictions_batch_format_export.py | 58 +++++------- e2e/test_predictions_batch_main_workflow.py | 18 ++-- ...st_predictions_batch_prediction_process.py | 33 +++---- laborious/activities/gates.py | 44 ++++++--- laborious/activities/mlflow.py | 5 +- laborious/metrics.py | 4 +- .../utils/models/minio_dataframe_payload.py | 3 +- laborious/worker/worker.py | 5 +- laborious/workflows/drift.py | 4 +- laborious/workflows/minimal_retrain.py | 4 +- laborious/workflows/predictions_batch.py | 8 +- laborious/workflows/simple_metrics.py | 2 +- .../sub_workflows/prediction_process.py | 2 - requirements.txt | 2 +- tests/laborious/activities/test_gates.py | 73 +++++++++++++-- tests/laborious/activities/test_mlflow.py | 2 + tests/laborious/activities/test_storage.py | 12 +-- .../models/test_minio_dataframe_payload.py | 14 ++- .../subworkflows/test_prediction_process.py | 47 +++++----- tests/laborious/workflows/test_drift.py | 26 +++--- .../workflows/test_minimal_retrain.py | 35 +++++-- .../workflows/test_predictions_batch.py | 40 +++++--- .../workflows/test_simple_metrics.py | 92 +++++++++---------- 24 files changed, 331 insertions(+), 234 deletions(-) diff --git a/e2e/conftest.py b/e2e/conftest.py index 61f86f2..5f3aa96 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -247,15 +247,17 @@ def mock_minio_repository(): def mock_pi_web_api_repository(): """Mock PI Web API repository for PI Web API operations.""" mock_repo = MagicMock() - mock_repo.write_value = AsyncMock( - return_value={ - 'Items': [ - { - 'WebId': 'web_id_1' - } - ] - } - ) + + async def _write_value(web_ids, value, metadata=None, **kwargs): + """ + Mirror successful PI writes: one response item per requested web_id. + + write_pi_web_api_data passes the list into process_pi_web_api_response (not a + wrapped {'Items': ...} envelope). + """ + return [{'WebId': wid, 'Errors': []} for wid in web_ids] + + mock_repo.write_value = AsyncMock(side_effect=_write_value) mock_repo.close = MagicMock() return mock_repo @@ -266,7 +268,7 @@ def mock_opc_repository(): mock_repo.write_data = AsyncMock( return_value=(True, {'response_time': 0.1}) ) - mock_repo.disconnect = MagicMock() + mock_repo.disconnect = AsyncMock() return mock_repo @pytest_asyncio.fixture @@ -279,7 +281,8 @@ def patch_create_engine(postgres_engine): @pytest_asyncio.fixture def patch_minio_repository(mock_minio_repository): """Patch MinioRepository to return mock.""" - with patch('sientia_do.repository.minio_repository.MinioRepository', return_value=mock_minio_repository): + # Patch where Activities resolves the symbol (import binds the original class). + with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository): yield @pytest_asyncio.fixture @@ -407,10 +410,13 @@ async def test_activities( 'password': 'test', }, minio_config={ - 'endpoint_url': 'http://localhost:9000', + # Host:port only; Minio() prepends http(s):// from the secure flag. + 'endpoint_url': 'localhost:9000', 'access_key': 'test', 'secret_key': 'test', 'default_bucket': 'test-bucket', + 'retention_hours': 24, + 'secure': False, }, opc_config={}, pi_web_api_config={ @@ -452,7 +458,6 @@ async def temporal_worker(temporal_test_env, test_activities): test_activities.load_custom_query, test_activities.load_query_with_minio_offload, test_activities.cleanup_minio_objects_expired, - test_activities.get_last_timestamp, test_activities.input_gate, test_activities.request_transform, test_activities.mlflow_response_gate, @@ -465,6 +470,7 @@ async def temporal_worker(temporal_test_env, test_activities): test_activities.write_pi_web_api_data, test_activities.write_opc_data, test_activities.export_data_to_postgres, + test_activities.export_payload_to_postgres, test_activities.write_metrics, ], ) as worker: diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py index e7bdd91..d7e7621 100644 --- a/e2e/test_predictions_batch_format_export.py +++ b/e2e/test_predictions_batch_format_export.py @@ -4,6 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios. import asyncio from datetime import datetime +from typing import Any, cast from unittest.mock import ANY, AsyncMock, patch, call import pandas as pd @@ -25,13 +26,13 @@ base_input_data = { 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -191,7 +192,6 @@ async def test_scenario_3_1_1_default_prediction_export( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0.5, }, - endpoint='test_endpoint', metadata={ 'model_id': 311, 'model_name': 'test_model', @@ -205,7 +205,6 @@ async def test_scenario_3_1_1_default_prediction_export( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 311, 'model_name': 'test_model', @@ -217,7 +216,8 @@ async def test_scenario_3_1_1_default_prediction_export( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ call('addr_1', 0.5, 'float', ANY, { @@ -300,7 +300,8 @@ async def test_scenario_3_1_2_export_with_opc_only( await start_and_await_workflow(client, input_data, workflow_id) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ call('addr_1', 0.5, 'float', ANY, { @@ -382,7 +383,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0.5, }, - endpoint='test_endpoint', metadata={ 'model_id': 313, 'model_name': 'test_model', @@ -396,7 +396,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 313, 'model_name': 'test_model', @@ -408,7 +407,8 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_not_called() + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_not_called() assert_prediction(postgres_engine, model_id) @@ -459,7 +459,8 @@ async def test_scenario_3_1_4_export_without_optional_outputs( await start_and_await_workflow(client, input_data, workflow_id) test_activities.pi_web_api_client.write_value.assert_not_called() - test_activities.opc_repository['1'].write_data.assert_not_called() + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_not_called() assert_prediction(postgres_engine, model_id) @@ -535,7 +536,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0.5, }, - endpoint='test_endpoint', metadata={ 'model_id': 315, 'model_name': 'test_model', @@ -549,7 +549,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( 'Timestamp': '2024-01-01 12:00:00+0000', 'Value': 0, }, - endpoint='test_endpoint', metadata={ 'model_id': 315, 'model_name': 'test_model', @@ -561,7 +560,8 @@ async def test_scenario_3_1_5_export_without_transformed_data( any_order=True, ) - test_activities.opc_repository['1'].write_data.assert_has_calls( + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.assert_has_calls( [ call('addr_1', 0.5, 'float', ANY, { @@ -695,7 +695,8 @@ async def test_scenario_3_2_2_opc_write_error( insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) print("[TEST] ✓ Data inserted successfully") - test_activities.opc_repository['1'].write_data.return_value = (False, { + opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) + opc_write_data.return_value = (False, { 'notification_id': 'OPC_WRITE_DATA_ERROR_1', 'message': 'OPC server unavailable', 'block': 'opc_repository', @@ -774,25 +775,14 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) print("[TEST] ✓ Data inserted successfully") - test_activities.pi_web_api_client.write_value = AsyncMock(side_effect=[ - { - 'Items': [ - { - 'WebId': 'web_id_1', - 'Errors': [], - }, - ] - }, - Exception('Tag write failed'), - { - 'Items': [ - { - 'WebId': 'web_id_2', - 'Errors': [], - }, - ] - }, - ]) + test_activities.pi_web_api_client.write_value = AsyncMock( + side_effect=[ + # Prediction batch: two web_ids requested, only one acknowledged. + [{'WebId': 'web_id_1', 'Errors': []}], + # Confidence write succeeds. + [{'WebId': 'web_id_2', 'Errors': []}], + ] + ) input_data = get_base_input_data(model_id) input_data['pi_web_api_output_config'] = { diff --git a/e2e/test_predictions_batch_main_workflow.py b/e2e/test_predictions_batch_main_workflow.py index f8469ef..ddac7ec 100644 --- a/e2e/test_predictions_batch_main_workflow.py +++ b/e2e/test_predictions_batch_main_workflow.py @@ -100,13 +100,13 @@ async def test_scenario_1_1_1_happy_path_complete_success( 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -238,13 +238,13 @@ async def test_scenario_1_2_1_sql_query_execution_error( 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -407,13 +407,13 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( 'table_name': 'predictions', 'transform_table_name': 'transformed_data', 'input_filters': { - 'EMPTY_DATA': {'policy': 'STOP', 'config': {}}, + 'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index d32980e..98f9429 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/e2e/test_predictions_batch_prediction_process.py @@ -5,6 +5,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios. import asyncio from datetime import datetime from decimal import Decimal +from typing import Any from unittest.mock import MagicMock, patch import pandas as pd @@ -27,15 +28,15 @@ base_input_data = { 'transform_table_name': 'transformed_data', 'input_filters': { 'SPECIFIC_VARIABLES_NULL_VALUES': { - 'policy': 'CONTINUE', # Continue despite issues, not STOP - 'config': {'variables': ['sensor_1']}, + 'POLICY': 'CONTINUE', # Continue despite issues, not STOP + 'CONFIG': {'variables': ['sensor_1']}, }, }, 'mlflow_transform_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'mlflow_predict_filters': { - 'API_ERROR': {'policy': 'STOP', 'config': {}}, + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, }, 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], 'opc_output_config': {}, @@ -59,7 +60,7 @@ def get_base_input_data(model_id): 'query': base_query.format(model_id=model_id), } -def insert_sample_data(postgres_engine, model_id, values: list[tuple]): +def insert_sample_data(postgres_engine, model_id, values: list[Any]): with postgres_engine.begin() as conn: conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}")) @@ -111,7 +112,7 @@ async def start_and_await_workflow(client, input_data, workflow_id): pytest.fail("Workflow execution timed out after 60 seconds") def assert_continue( - postgres_engine, model_id, prediction_confidence: Decimal = 2, + postgres_engine, model_id, prediction_confidence: Decimal = Decimal(2), comments: str = 'Input data with bad quality', ): print("\n[TEST] 4. Verifying prediction was created despite warnings...") @@ -236,7 +237,7 @@ async def test_scenario_2_1_2_input_gate_triggers_stop( print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) - input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'STOP' + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP' print("\n[TEST] 2. Starting workflow that should stop at input gate...") workflow_id = f'test-input-stop-{datetime.now().timestamp()}' @@ -285,7 +286,7 @@ async def test_scenario_2_1_3_input_gate_triggers_repeat( print("[TEST] ✓ Data and previous prediction inserted") input_data = get_base_input_data(model_id) - input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['policy'] = 'REPEAT' + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' print("\n[TEST] 2. Starting workflow that should trigger REPEAT...") workflow_id = f'test-input-repeat-{datetime.now().timestamp()}' @@ -345,7 +346,7 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue( insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'CONTINUE' + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE' print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...") workflow_id = f'test-transform-continue-{datetime.now().timestamp()}' @@ -355,7 +356,7 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), - comments='Bad data model', + comments='Unknown MLFlow API error', ) print("\n[TEST] ✓ All assertions passed!") @@ -397,7 +398,7 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop( print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'STOP' + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP' print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...") workflow_id = f'test-transform-stop-{datetime.now().timestamp()}' @@ -449,7 +450,7 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat( print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) - input_data['mlflow_transform_filters']['API_ERROR']['policy'] = 'REPEAT' + input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT' print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...") workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}' @@ -516,7 +517,7 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue( insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'CONTINUE' + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE' print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...") workflow_id = f'test-predict-continue-{datetime.now().timestamp()}' @@ -526,7 +527,7 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), - comments='Bad predict model', + comments='Unknown MLFlow API error', ) print("\n[TEST] ✓ All assertions passed!") @@ -567,7 +568,7 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop( print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'STOP' + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP' print("\n[TEST] 2. Starting workflow that should stop at predict gate...") workflow_id = f'test-predict-stop-{datetime.now().timestamp()}' @@ -617,7 +618,7 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat( print("[TEST] ✓ Data and previous prediction inserted") input_data = get_base_input_data(model_id) - input_data['mlflow_predict_filters']['API_ERROR']['policy'] = 'REPEAT' + input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT' input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] # REPEAT first print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...") diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index cac6c38..1d561c5 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -118,6 +118,22 @@ class Gates(MinioManager): def __del__(self): self.close() + @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]: """ @@ -171,10 +187,11 @@ class Gates(MinioManager): 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, config['config']): + if input_filter_functions[fil](data, filter_config): self.debug(f'Data not passed the input filter {fil}:{config}', metadata) - filter_output.append(config['policy']) + filter_output.append(policy) except Exception as e: trace = traceback.format_exc() await self.send_notification_async( @@ -235,8 +252,10 @@ class Gates(MinioManager): 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'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) @@ -254,17 +273,18 @@ class Gates(MinioManager): 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, config): - filter_output.append(config['policy']) - comments.append(status['message']) + 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=data['content']['message'], + message=status.get('message', 'Unknown MLFlow API error'), block='mlflow_gate', level=NotificationLevel.ERROR, - attachment_content=data['content']['traceback'], + attachment_content=status.get('traceback'), ) except Exception as e: trace = traceback.format_exc() @@ -341,9 +361,10 @@ class Gates(MinioManager): for fil, config in filters.items(): if fil not in mlflow_content_filter_functions: continue + policy, filter_config = self._read_filter_entry(config) try: - if mlflow_content_filter_functions[fil](data, config): - filter_output.append(config['policy']) + if mlflow_content_filter_functions[fil](data, filter_config): + filter_output.append(policy) await self.send_notification_async( metadata=metadata, notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}', @@ -703,7 +724,6 @@ class Gates(MinioManager): self.info(f'Writing metrics for model {metadata["model_name"]}', metadata) - core_tags = { 'pod_id': self.pod_id, 'runtime': self.runtime, diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index ea97f3b..ad9b6ae 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -1,7 +1,5 @@ from temporalio import activity, workflow - - with workflow.unsafe.imports_passed_through(): import traceback from typing import Any @@ -20,8 +18,9 @@ with workflow.unsafe.imports_passed_through(): now, ) from sientia_do.utils.formatters import create_sample_dict - from laborious.utils.repository.minio_manager import MinioManager + 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 diff --git a/laborious/metrics.py b/laborious/metrics.py index a67e574..776b826 100644 --- a/laborious/metrics.py +++ b/laborious/metrics.py @@ -25,9 +25,7 @@ Metric Labels: """ from prometheus_client import Counter, Gauge, Histogram -from sientia_do.observability.metrics import ( - CORE_LABELS -) +from sientia_do.observability.metrics import CORE_LABELS # Application health metric APP_UP = Gauge( diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 42b60f4..bb3f2fd 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -162,7 +162,7 @@ class MinioDataFramePayload: """ Return True if the payload has some data internally or in MinIO. """ - return (self.data is not None and not self.data != {}) or self.object_key is not None + return (self.data is not None and self.data != {}) or self.object_key is not None @classmethod async def from_dataframe( @@ -200,7 +200,6 @@ class MinioDataFramePayload: data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status ) - if last_timestamp is None: last_timestamp = max(dataframe['timestamp'].values.tolist()) diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 5202f2d..fa09259 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -219,20 +219,21 @@ async def main(): logger.custom_info('Workers started successfully', metadata) + 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() - # Exit with a non-zero status code to indicate failure to Kubernetes metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN - sys.exit(1) + sys.exit(exit_code) def start_prometheus_server(): diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index b98fc56..2faff32 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -46,7 +46,7 @@ class Drift: ORDER BY timestamp ASC """ # nosec B608 - values come from internal Temporal workflow config, not user input - target_data_handler = workflow.start_local_activity_method( + target_data_handler = workflow.start_activity_method( Activities.load_custom_query, { **metadata, @@ -58,7 +58,7 @@ class Drift: start_to_close_timeout=timedelta(seconds=300), ) - reference_data_handler = workflow.start_local_activity_method( + reference_data_handler = workflow.start_activity_method( Activities.get_reference_data, {**metadata, 'model_name': input_data['model_name']}, retry_policy=retry_policy, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index f0cf0e8..f8f5151 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -7,6 +7,7 @@ with workflow.unsafe.imports_passed_through(): 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') @@ -84,7 +85,8 @@ class MinimalRetrain: start_to_close_timeout=timedelta(seconds=600), ) - if not storage_result.has_data(): + 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( diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 427a470..36e2e97 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -105,12 +105,14 @@ class PredictionsBatch: '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'}}), + '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'}} + 'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}} ), 'mlflow_predict_filters': input_data.get( - 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + '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']), diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index 20b60ab..38b47c4 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -47,7 +47,7 @@ class SimpleMetrics: p."timestamp" desc; """ # nosec B608 - values come from internal Temporal workflow config, not user input - target_data = await workflow.execute_local_activity_method( + target_data = await workflow.execute_activity_method( Activities.load_custom_query, { **metadata, diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index fbac978..3493b59 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -7,7 +7,6 @@ with workflow.unsafe.imports_passed_through(): 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='subworkflow.prediction_process') @@ -112,7 +111,6 @@ class PredictionProcess: start_to_close_timeout=timedelta(minutes=5), ) raise e - async def _run_prediction_pipeline( self, diff --git a/requirements.txt b/requirements.txt index dd4fa1a..f0eedf8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.3 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4 git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.7 prometheus-client botocore diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 367f5b4..3d072b9 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -90,7 +90,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac ) input_data = { **metadata, - 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -103,7 +103,7 @@ async def test_input_gate_filter_exception(mock_input_filter_functions, gates_ac gates_activity.send_notification_async.assert_called_once_with( metadata=metadata['metadata'], notification_id='INTPUT_GATE_ERROR__EMPTY_DATA', - message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error", + message="Error in filter EMPTY_DATA:{'POLICY': 'STOP', 'CONFIG': {}}: \n Test error", block='input_gate', level=NotificationLevel.ERROR, attachment_content=ANY, @@ -133,7 +133,7 @@ async def test_input_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, 'data': _minio_payload(DataFrame({'value': []})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -147,11 +147,45 @@ async def test_input_gate_with_filter(gates_activity): @mark.asyncio -async def test_input_gate_with_filter_not_caught(gates_activity): +async def test_input_gate_with_filter_lowercase_keys(gates_activity): # Arrange input_data = { **metadata, 'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}}, + 'data': _minio_payload(DataFrame({'value': []})), + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.input_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'Input data with bad quality') + + +@mark.asyncio +async def test_input_gate_with_filter_capitalized_keys(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'EMPTY_DATA': {'Policy': 'STOP', 'Config': {}}}, + 'data': _minio_payload(DataFrame({'value': []})), + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.input_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'Input data with bad quality') + + +@mark.asyncio +async def test_input_gate_with_filter_not_caught(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], } @@ -248,7 +282,7 @@ async def test_mlflow_response_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'API_ERROR': {'policy': 'STOP'}}, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, 'data': _minio_payload( {'content': {'message': 'API error occurred', 'traceback': 'error trace'}}, status={'success': False, 'message': 'API error occurred'}, @@ -266,12 +300,33 @@ async def test_mlflow_response_gate_with_filter(gates_activity): gates_activity.send_notification_async.assert_called() +@mark.asyncio +async def test_mlflow_response_gate_with_filter_capitalized_keys(gates_activity): + # Arrange + input_data = { + **metadata, + 'filters': {'API_ERROR': {'Policy': 'STOP'}}, + 'data': _minio_payload( + {'content': {'message': 'API error occurred', 'traceback': 'error trace'}}, + status={'success': False, 'message': 'API error occurred'}, + ), + 'type': 'test', + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + } + + # Act + result = await gates_activity.mlflow_response_gate(input_data) + + # Assert + assert result == ('STOP', -1, 'API error occurred') + + @mark.asyncio async def test_mlflow_response_gate_with_filter_not_caught(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'API_ERROR': {'policy': 'STOP'}}, + 'filters': {'API_ERROR': {'POLICY': 'STOP'}}, 'data': _minio_payload( {'content': {'message': 'success'}}, status={'success': True}, @@ -364,7 +419,7 @@ async def test_mlflow_content_gate_with_filter(gates_activity): # Arrange input_data = { **metadata, - 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, + 'filters': {'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}}, 'data': _minio_payload(DataFrame({'value': [None, None, None]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], @@ -402,7 +457,7 @@ async def test_mlflow_content_gate_with_filter_not_caught(gates_activity): async def test_mlflow_content_gate_filter_returns_false(gates_activity): input_data = { **metadata, - 'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}}, + 'filters': {'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}}, 'data': _minio_payload(DataFrame({'value': [1, 2, 3]})), 'type': 'test', 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], @@ -823,6 +878,7 @@ async def test_write_metrics(mock_metrics, gates_activity): core_tags = { 'pod_id': gates_activity.pod_id, 'runtime': gates_activity.runtime, + 'operation_type': 'predict', 'model_name': metadata['metadata']['model_name'], 'workflow_name': metadata['metadata']['workflow_name'], } @@ -930,6 +986,7 @@ async def test_write_metrics_with_none_opc_response_time(mock_metrics, gates_act tags={ 'pod_id': gates_activity.pod_id, 'runtime': gates_activity.runtime, + 'operation_type': 'predict', 'model_name': metadata['metadata']['model_name'], 'workflow_name': metadata['metadata']['workflow_name'], 'opc_server_id': 'server1', diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index e495be8..daaebab 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -180,6 +180,7 @@ async def test_request_transform_failure(mock_from_dataframe, mlflow): operation='transform', status=transform_response, workflow_metadata=metadata['metadata'], + last_timestamp=payload.last_timestamp, ) assert response_data == mock_from_dataframe.return_value @@ -250,6 +251,7 @@ async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, ml operation='predict', status=predict_response, workflow_metadata=metadata['metadata'], + last_timestamp=payload.last_timestamp, ) assert response_data == mock_from_dataframe.return_value diff --git a/tests/laborious/activities/test_storage.py b/tests/laborious/activities/test_storage.py index 9e76c5f..ce99ff7 100644 --- a/tests/laborious/activities/test_storage.py +++ b/tests/laborious/activities/test_storage.py @@ -204,9 +204,7 @@ async def test_cleanup_minio_objects_expired(mock_now, storage): data_mock = MagicMock() data_mock.cleanup_prefix.return_value = 'training_datasets/m' - result = await storage.cleanup_minio_objects_expired( - {**metadata, 'data': data_mock} - ) + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) assert result['deleted_count'] == 1 assert result['failed_count'] == 0 @@ -291,9 +289,7 @@ async def test_cleanup_minio_objects_expired_delete_fails(mock_now, storage): data_mock = MagicMock() data_mock.cleanup_prefix.return_value = 'training_datasets/m' - result = await storage.cleanup_minio_objects_expired( - {**metadata, 'data': data_mock} - ) + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) assert result['deleted_count'] == 0 assert result['failed_count'] == 1 @@ -312,9 +308,7 @@ async def test_cleanup_minio_objects_expired_list_objects_error(mock_now, storag data_mock = MagicMock() data_mock.cleanup_prefix.return_value = 'training_datasets/m' - result = await storage.cleanup_minio_objects_expired( - {**metadata, 'data': data_mock} - ) + result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock}) assert result['deleted_count'] == 0 assert result['failed_count'] == 0 diff --git a/tests/laborious/utils/models/test_minio_dataframe_payload.py b/tests/laborious/utils/models/test_minio_dataframe_payload.py index de1a334..c27765d 100644 --- a/tests/laborious/utils/models/test_minio_dataframe_payload.py +++ b/tests/laborious/utils/models/test_minio_dataframe_payload.py @@ -90,22 +90,20 @@ async def test_retrieve_downloads_parquet_when_offloaded(): def test_build_object_key(): key, prefix = _build_object_key('my-model', 'initial', '2024-01-01_00-00-00') - assert key == 'training_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet' - assert prefix == 'training_datasets/my-model' + assert key == 'prediction_datasets/my-model/my-model-initial-2024-01-01_00-00-00.parquet' + assert prefix == 'prediction_datasets/my-model' def test_build_object_key_strips_slashes(): key, prefix = _build_object_key(' /my-model/ ', 'transform', '2024-06-15_10-30-45') - assert prefix == 'training_datasets/my-model' - assert key.startswith('training_datasets/my-model/') + assert prefix == 'prediction_datasets/my-model' + assert key.startswith('prediction_datasets/my-model/') def test_estimate_size_bytes_fallback(): df = DataFrame({'a': [1, 2]}) - original_to_dict = df.to_dict - df.to_dict = lambda *a, **kw: (_ for _ in ()).throw(RuntimeError('to_dict failed')) - size = MinioDataFramePayload.estimate_size_bytes(df) - df.to_dict = original_to_dict + with patch.object(df, 'to_dict', side_effect=RuntimeError('to_dict failed')): + size = MinioDataFramePayload.estimate_size_bytes(df) assert isinstance(size, int) assert size > 0 diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index db75b8c..60f3744 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -6,15 +6,6 @@ from laborious.activities.activities import Activities from laborious.workflows.sub_workflows.prediction_process import PredictionProcess -@fixture(autouse=True) -def _passthrough_from_dict(): - with patch( - 'laborious.workflows.sub_workflows.prediction_process.MinioDataFramePayload.from_dict', - side_effect=lambda x: x, - ): - yield - - @fixture def prediction_process(): return PredictionProcess() @@ -38,7 +29,9 @@ async def test_run(workflow_mock, prediction_process): # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, @@ -128,7 +121,7 @@ async def test_run(workflow_mock, prediction_process): { **metadata, 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], }, @@ -143,7 +136,7 @@ async def test_run(workflow_mock, prediction_process): Activities.request_predict, { **metadata, - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], }, @@ -174,8 +167,8 @@ async def test_run(workflow_mock, prediction_process): { 'metadata': metadata, 'path_flag': 'continue', - 'data': 'predicted_data', - 'transformed_data': 'transformed_data', + 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + 'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'prediction_confidence': 0.95, 'timestamp': '2024-01-01', 'model_id': 1, @@ -199,7 +192,9 @@ async def test_run_stop_at_input_gate(workflow_mock, prediction_process): # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, @@ -251,7 +246,9 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, @@ -336,7 +333,9 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, @@ -421,7 +420,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], **metadata, @@ -441,7 +440,9 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p # Arrange data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, @@ -527,7 +528,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p Activities.mlflow_content_gate, { 'filters': input_data['mlflow_transform_filters'], - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'type': 'transform', 'path_priority': input_data['path_priority'], **metadata, @@ -542,7 +543,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p call( Activities.request_predict, { - 'data': 'transformed_data', + 'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, 'model_name': input_data['model_name'], 'model_config': input_data['model_config'], **metadata, @@ -784,7 +785,9 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): data_payload = MagicMock() data_payload.cleanup_prefix.return_value = 'training_datasets/test' - data_payload.__getitem__ = lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + data_payload.__getitem__ = ( + lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock() + ) input_data = { 'metadata': metadata, 'data': data_payload, diff --git a/tests/laborious/workflows/test_drift.py b/tests/laborious/workflows/test_drift.py index 8f013d3..c31dafc 100644 --- a/tests/laborious/workflows/test_drift.py +++ b/tests/laborious/workflows/test_drift.py @@ -45,10 +45,9 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) + workflow_mock.execute_activity_method = AsyncMock(return_value=None) # Act await drift.run(input_data) @@ -64,7 +63,7 @@ async def test_run(workflow_mock: AsyncMock, drift: Drift): ORDER BY timestamp ASC """ - workflow_mock.start_local_activity_method.assert_has_calls( + workflow_mock.start_activity_method.assert_has_calls( [ call( Activities.load_custom_query, @@ -140,16 +139,15 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, drift: Drift): target_data = None reference_data = {'data': 'test_reference_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] - workflow_mock.execute_local_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock() workflow_mock.execute_activity_method = AsyncMock() # Act await drift.run(input_data) # Assert - Should not call calculate_drift or export - workflow_mock.execute_local_activity_method.assert_not_called() workflow_mock.execute_activity_method.assert_not_called() @@ -175,9 +173,8 @@ async def test_run_empty_drift_data(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = None - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) workflow_mock.execute_activity_method = AsyncMock() # Act @@ -226,10 +223,9 @@ async def test_run_default_chunk_period(workflow_mock: AsyncMock, drift: Drift): reference_data = {'data': 'test_reference_data'} drift_data = {'drift': 'test_drift_data'} - workflow_mock.start_local_activity_method.side_effect = [target_data, reference_data] - - workflow_mock.execute_local_activity_method.return_value = drift_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.start_activity_method.side_effect = [target_data, reference_data] + workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data) + workflow_mock.execute_activity_method = AsyncMock(return_value=None) # Act await drift.run(input_data) diff --git a/tests/laborious/workflows/test_minimal_retrain.py b/tests/laborious/workflows/test_minimal_retrain.py index be49c89..94d86da 100644 --- a/tests/laborious/workflows/test_minimal_retrain.py +++ b/tests/laborious/workflows/test_minimal_retrain.py @@ -1,4 +1,4 @@ -from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +from unittest.mock import ANY, AsyncMock, call, patch from pytest import fixture, mark @@ -39,8 +39,15 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain): }, } - storage_result = MagicMock() - storage_result.has_data.return_value = True + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } workflow_mock.execute_activity_method = AsyncMock( side_effect=[ @@ -163,8 +170,15 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim }, } - storage_result = MagicMock() - storage_result.has_data.return_value = False + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } workflow_mock.execute_activity_method = AsyncMock( side_effect=[ @@ -218,8 +232,15 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim }, } - storage_result = MagicMock() - storage_result.has_data.return_value = True + storage_result = { + 'last_timestamp': '2024-01-01 00:00:00+0000', + 'status': {'success': True}, + 'data': {'timestamp': {0: '2024-01-01 00:00:00+0000'}, 'value': {0: 1.0}}, + 'bucket': None, + 'object_key': None, + 'object_prefix': None, + 'uri': None, + } workflow_mock.execute_activity_method = AsyncMock( side_effect=[ diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 279499f..de52ad8 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -12,12 +12,10 @@ def predictions_batch() -> PredictionsBatch: metadata = { - 'metadata': { - 'model_id': 'test_model_id', - 'model_name': 'test_model', - 'workflow_name': 'predictions_batch', - 'schedule_name': 'test_schedule', - }, + 'model_id': 'test_model_id', + 'model_name': 'test_model', + 'workflow_name': 'predictions_batch', + 'schedule_name': 'test_schedule', } @@ -49,7 +47,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch call( Activities.load_query_with_minio_offload, { - **metadata, + 'metadata': metadata, 'query': input_data['query'], 'datetime_columns': input_data.get('datetime_columns', []), 'model_name': input_data['model_name'], @@ -60,19 +58,39 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch ] ) prediction_input = { - 'metadata': metadata, + 'metadata': {'metadata': metadata}, 'data': activity_return, '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'}}), + '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'}} + 'mlflow_transform_filters', + { + 'API_ERROR': { + 'POLICY': 'STOP', + 'CONFIG': {}, + } + }, ), 'mlflow_predict_filters': input_data.get( - 'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}} + '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']), diff --git a/tests/laborious/workflows/test_simple_metrics.py b/tests/laborious/workflows/test_simple_metrics.py index 0d96bc5..dbda735 100644 --- a/tests/laborious/workflows/test_simple_metrics.py +++ b/tests/laborious/workflows/test_simple_metrics.py @@ -42,9 +42,8 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None]) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) @@ -66,7 +65,7 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): p."timestamp" desc; """ - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.load_custom_query, @@ -80,29 +79,30 @@ async def test_run(workflow_mock: AsyncMock, simple_metrics: SimpleMetrics): start_to_close_timeout=ANY, ), call( - Activities.calculate_simple_metrics, + Activities.export_data_to_postgres, { **metadata, - 'model_id': input_data['model_id'], - 'target_data': target_data, - 'metrics': input_data['metrics'], - 'interval_minutes': input_data['interval_minutes'], + 'data': simple_metrics_data, + 'schema': input_data['schema'], + 'table_name': input_data['target_table_name'], + 'timestamp_conversion': { + 'column': 'timestamp', + 'format': DATETIME_FORMAT_WITH_TZ, + }, }, retry_policy=ANY, start_to_close_timeout=ANY, ), ] ) - - # Assert - Check export_data_to_postgres call - workflow_mock.execute_activity_method.assert_called_once_with( - Activities.export_data_to_postgres, + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_simple_metrics, { **metadata, - 'data': simple_metrics_data, - 'schema': input_data['schema'], - 'table_name': input_data['target_table_name'], - 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': input_data['metrics'], + 'interval_minutes': input_data['interval_minutes'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -128,15 +128,13 @@ async def test_run_empty_target_data(workflow_mock: AsyncMock, simple_metrics: S target_data = None - workflow_mock.execute_local_activity_method.return_value = target_data - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(return_value=target_data) # Act await simple_metrics.run(input_data) # Assert - Should not call calculate_simple_metrics or export - assert workflow_mock.execute_local_activity_method.call_count == 1 - workflow_mock.execute_activity_method.assert_not_called() + assert workflow_mock.execute_activity_method.call_count == 1 @mark.asyncio @@ -159,16 +157,15 @@ async def test_run_empty_simple_metrics(workflow_mock: AsyncMock, simple_metrics target_data = {'data': 'test_target_data'} simple_metrics_data = None - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(return_value=target_data) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) # Assert - Should call calculate_simple_metrics but not export - assert workflow_mock.execute_local_activity_method.call_count == 2 - workflow_mock.execute_activity_method.assert_not_called() + workflow_mock.execute_activity_method.assert_called_once() + workflow_mock.execute_local_activity_method.assert_called_once() @mark.asyncio @@ -191,33 +188,28 @@ async def test_run_default_metrics(workflow_mock: AsyncMock, simple_metrics: Sim target_data = {'data': 'test_target_data'} simple_metrics_data = {'metrics': 'test_simple_metrics_data'} - workflow_mock.execute_local_activity_method.side_effect = [target_data, simple_metrics_data] - - workflow_mock.execute_activity_method = AsyncMock() + workflow_mock.execute_activity_method = AsyncMock(side_effect=[target_data, None]) + workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data) # Act await simple_metrics.run(input_data) # Assert - Check calculate_simple_metrics call with default metrics - workflow_mock.execute_local_activity_method.assert_has_calls( - [ - call( - Activities.load_custom_query, - ANY, - retry_policy=ANY, - start_to_close_timeout=ANY, - ), - call( - Activities.calculate_simple_metrics, - { - **metadata, - 'model_id': input_data['model_id'], - 'target_data': target_data, - 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value - 'interval_minutes': input_data['interval_minutes'], - }, - retry_policy=ANY, - start_to_close_timeout=ANY, - ), - ] + workflow_mock.execute_activity_method.assert_any_call( + Activities.load_custom_query, + ANY, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + workflow_mock.execute_local_activity_method.assert_called_once_with( + Activities.calculate_simple_metrics, + { + **metadata, + 'model_id': input_data['model_id'], + 'target_data': target_data, + 'metrics': ['rmse', 'mse', 'mae', 'r2'], # Default value + 'interval_minutes': input_data['interval_minutes'], + }, + retry_policy=ANY, + start_to_close_timeout=ANY, ) From 0e3ec6463f38f03418c9ee8a2fa8dacc44dddc43 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 09:44:35 -0300 Subject: [PATCH 23/51] SIENTIAPDE-1712 Enhance E2E testing with MinIO support and update documentation - Updated `requirements-dev.txt` to include MinIO support in testcontainers. - Added a new fixture for MinIO container setup in `conftest.py` to facilitate E2E tests involving S3-compatible storage. - Introduced a new test fixture for activities using a real MinIO container in `conftest.py`. - Updated E2E test scenarios and documentation to reflect the integration of MinIO for offload uploads and clarified error handling in workflows. - Refactored existing tests to improve clarity and maintainability. --- e2e/conftest.py | 140 ++++- e2e/helpers.py | 174 +++++ e2e/scenarios.md | 39 +- e2e/test_child_workflows_e2e.py | 74 +++ e2e/test_minio_offload.py | 124 ++++ e2e/test_predictions_batch_format_export.py | 304 +++------ e2e/test_predictions_batch_main_workflow.py | 289 +++------ ...st_predictions_batch_prediction_process.py | 595 +++++------------- laborious/activities/api.py | 9 +- laborious/activities/storage.py | 13 +- requirements-dev.txt | 2 +- 11 files changed, 870 insertions(+), 893 deletions(-) create mode 100644 e2e/helpers.py create mode 100644 e2e/test_child_workflows_e2e.py create mode 100644 e2e/test_minio_offload.py diff --git a/e2e/conftest.py b/e2e/conftest.py index 5f3aa96..7825335 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -9,6 +9,7 @@ import pandas as pd import pytest import pytest_asyncio from sqlalchemy import create_engine, text +from testcontainers.minio import MinioContainer from testcontainers.postgres import PostgresContainer from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker @@ -26,6 +27,17 @@ TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017' TEST_DATABASE_NAME = 'test_db' +@pytest_asyncio.fixture(scope='session') +def minio_container(): + """ + MinIO S3-compatible storage for E2E tests that exercise real offload uploads. + """ + minio = MinioContainer() + minio.start() + yield minio + minio.stop() + + @pytest_asyncio.fixture(scope='session') def postgres_container(): """ @@ -187,6 +199,20 @@ def mock_mongo_client(): return mock_client +@pytest.fixture +def notification_inserts(mock_mongo_client): + """ + Mongo insert_one mock used by CoreNotificationHandler for notification persistence. + + Yields: + MagicMock for insert_one, reset before each test. + """ + mock_db = mock_mongo_client.__getitem__.return_value + mock_collection = mock_db.__getitem__.return_value + mock_collection.insert_one.reset_mock() + yield mock_collection.insert_one + + @pytest_asyncio.fixture def notification_handler(mock_logger, mock_mongo_client): """ @@ -439,6 +465,88 @@ async def test_activities( await activities.shutdown() +@pytest_asyncio.fixture(scope='function') +async def test_activities_real_minio( + postgres_engine, + postgres_container, + minio_container, + mock_logger, + notification_handler, + metrics_controller, + patch_create_engine, + patch_mlflow, + patch_pi_web_api_repository, + mock_opc_repository, +): + """ + Activities with a real MinIO testcontainer (no MinioRepository patch) for offload tests. + """ + minio_client = minio_container.get_client() + if not minio_client.bucket_exists('test-bucket'): + minio_client.make_bucket('test-bucket') + minio_port = minio_container.get_exposed_port(9000) + activities = Activities( + postgres_config={ + 'host': 'localhost', + 'port': postgres_container.get_exposed_port(5432), + 'user': 'test', + 'password': 'test', + 'dbname': 'test', + 'min_connections': 1, + 'max_connections': 5, + }, + mlflow_config={ + 'host': 'http://localhost', + 'port': '5000', + 'username': 'test', + 'password': 'test', + }, + minio_config={ + 'endpoint_url': f'localhost:{minio_port}', + 'access_key': 'minioadmin', + 'secret_key': 'minioadmin', + 'default_bucket': 'test-bucket', + 'retention_hours': 24, + 'secure': False, + }, + opc_config={}, + pi_web_api_config={ + 'base_url': 'http://localhost:8080', + 'auth_type': 'bearer', + 'auth_token': 'test_token', + }, + logger=mock_logger, + notification_handler=notification_handler, + ) + activities.opc_repository = {'1': mock_opc_repository} + try: + yield activities + finally: + await activities.shutdown() + + +def _worker_activity_list(test_activities: Activities): + return [ + test_activities.load_custom_query, + test_activities.load_query_with_minio_offload, + test_activities.cleanup_minio_objects_expired, + test_activities.input_gate, + test_activities.request_transform, + test_activities.mlflow_response_gate, + test_activities.mlflow_content_gate, + test_activities.request_predict, + test_activities.repeat_last_prediction, + test_activities.format_prediction, + test_activities.format_transformed_data, + test_activities.format_default_prediction, + test_activities.write_pi_web_api_data, + test_activities.write_opc_data, + test_activities.export_data_to_postgres, + test_activities.export_payload_to_postgres, + test_activities.write_metrics, + ] + + @pytest_asyncio.fixture(scope='function') async def temporal_test_env(): """Create Temporal test environment.""" @@ -454,24 +562,18 @@ async def temporal_worker(temporal_test_env, test_activities): temporal_test_env.client, task_queue='test-queue', workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], - activities=[ - test_activities.load_custom_query, - test_activities.load_query_with_minio_offload, - test_activities.cleanup_minio_objects_expired, - test_activities.input_gate, - test_activities.request_transform, - test_activities.mlflow_response_gate, - test_activities.mlflow_content_gate, - test_activities.request_predict, - test_activities.repeat_last_prediction, - test_activities.format_prediction, - test_activities.format_transformed_data, - test_activities.format_default_prediction, - test_activities.write_pi_web_api_data, - test_activities.write_opc_data, - test_activities.export_data_to_postgres, - test_activities.export_payload_to_postgres, - test_activities.write_metrics, - ], + activities=_worker_activity_list(test_activities), + ) as worker: + yield worker + + +@pytest_asyncio.fixture(scope='function') +async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio): + """Temporal worker backed by Activities using real MinIO testcontainer.""" + async with Worker( + temporal_test_env.client, + task_queue='test-queue', + workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction], + activities=_worker_activity_list(test_activities_real_minio), ) as worker: yield worker diff --git a/e2e/helpers.py b/e2e/helpers.py new file mode 100644 index 0000000..1aaa4f2 --- /dev/null +++ b/e2e/helpers.py @@ -0,0 +1,174 @@ +""" +Shared helpers for E2E tests (Temporal workflows + PostgreSQL). +""" + +import asyncio +from datetime import datetime +from decimal import Decimal +from typing import Any + +from sqlalchemy import text +from sqlalchemy.engine import Engine + + +async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0): + """ + Start a workflow and wait for its result. + + Args: + client: Temporal client from WorkflowEnvironment. + workflow_run: Workflow run method (e.g. PredictionsBatch.run). + input_data: Workflow input payload. + workflow_id: Unique workflow id. + timeout: Max seconds to wait for completion. + + Return: + Workflow result value. + """ + handle = await client.start_workflow( + workflow_run, + input_data, + id=workflow_id, + task_queue='test-queue', + ) + return await asyncio.wait_for(handle.result(), timeout=timeout) + + +def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any]) -> None: + """ + Replace laborious_data rows for a model_id with one row per value (sensor_1..n). + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Model id column value. + values: Per-sensor values; use string 'NULL' for SQL NULL. + """ + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + values_sql = [] + for i, value in enumerate(values): + values_sql.append(f""" + ({model_id}, 'sensor_{i + 1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') + """) + insert_sql = f""" + INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) + VALUES + {', '.join(values_sql)} + """ + conn.execute(text(insert_sql)) + + +def assert_prediction( + postgres_engine: Engine, + model_id: int, + prediction: float = 0.5, + prediction_confidence: int | Decimal = 0, + prediction_status: str = 'Good', + comments: str = '', +) -> None: + """ + Assert exactly one prediction row exists for model_id with expected columns. + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Expected model_id. + prediction: Expected prediction value. + prediction_confidence: Expected confidence (int or Decimal for numeric column). + prediction_status: Expected status string. + comments: Expected comments string. + """ + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id} ' + f'ORDER BY created_at ASC' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}' + row = prediction_rows[0] + assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}' + assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), ( + f'Expected prediction={prediction}, got {row[1]}' + ) + assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal( + str(prediction_confidence) + ), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}' + assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}" + assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + + +def assert_continue( + postgres_engine: Engine, + model_id: int, + prediction_confidence: Decimal = Decimal(2), + comments: str = 'Input data with bad quality', +) -> None: + """Assert one default-style prediction row after CONTINUE gate path.""" + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id}' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 1, 'Expected one prediction record despite warnings' + row = prediction_rows[0] + assert row[1] == 0, f'Expected prediction=0, got {row[1]}' + assert row[2] == prediction_confidence, ( + f'Expected prediction_confidence={prediction_confidence}, got {row[2]}' + ) + assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}" + assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + + +def assert_stop(postgres_engine: Engine, model_id: int) -> None: + """Assert no prediction rows for model_id.""" + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) + count = result_query.scalar() + assert count == 0, f'Expected no predictions, but found {count} records' + + +def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple) -> None: + """ + Assert two prediction rows for model_id both match last_prediction. + + Rows are compared in created_at order for stability. + + Args: + postgres_engine: SQLAlchemy engine. + model_id: Model id. + last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows. + """ + import pytest + + with postgres_engine.connect() as conn: + result_query = conn.execute( + text( + f'SELECT model_id, prediction, prediction_confidence, prediction_status ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id} ' + f'ORDER BY created_at ASC' + ) + ) + prediction_rows = result_query.fetchall() + assert len(prediction_rows) == 2, 'Expected two prediction records' + assert prediction_rows[0] == last_prediction, ( + f'Expected first row {last_prediction}, got {prediction_rows[0]}' + ) + assert prediction_rows[1] == last_prediction, ( + f'Expected second row {last_prediction}, got {prediction_rows[1]}' + ) + + +def make_workflow_id(prefix: str) -> str: + """Build a unique workflow id using a prefix and current timestamp.""" + return f'{prefix}-{datetime.now().timestamp()}' diff --git a/e2e/scenarios.md b/e2e/scenarios.md index eb9c606..5732cd3 100644 --- a/e2e/scenarios.md +++ b/e2e/scenarios.md @@ -2,6 +2,13 @@ This document describes all possible test scenarios for the `predictions_batch` workflow and its child workflows `prediction_process` and `format_and_export_prediction`. +## Running automated E2E tests (`e2e/`) + +- **Runtime**: Docker (or a Docker-compatible daemon) must be available so [testcontainers](https://testcontainers.com/) can start **PostgreSQL** and **MinIO** containers. +- **Dependencies**: install dev requirements (includes `testcontainers[postgres,minio]`). +- **Invocation**: run only integration-marked tests, for example: `pytest e2e/ -m integration`. +- **MinIO tests**: `e2e/test_minio_offload.py` exercises real S3 uploads; other E2E modules continue to mock MinIO on the worker used by most scenarios. + ## Workflow Overview The `predictions_batch` workflow: @@ -445,6 +452,8 @@ The `predictions_batch` workflow: ### 3.2 Error Scenarios +These paths do **not** rely on Temporal activity retries for export failures: the write activities run once, errors are handled inside the activity, and the **workflow completes successfully** with degraded metadata on the persisted prediction (`prediction_confidence` and `comments`). + #### Scenario 3.2.1: PI Web API Write Error **Description**: PI Web API export fails @@ -453,15 +462,16 @@ The `predictions_batch` workflow: - PI Web API service unavailable or invalid config **Expected Behavior**: -- `write_pi_web_api_data` raises exception -- Notification sent -- Workflow fails after retries -- PostgreSQL export may not execute (depends on execution order) +- `write_pi_web_api_data` surfaces the failure (exception handled in the activity layer) +- Notification may be sent +- Workflow **completes** (does not fail) +- Prediction row is still written to PostgreSQL with error confidence **13** and a comment describing the PI error +- Subsequent steps (e.g. OPC, Postgres) still run per workflow order with the updated prediction payload **Assertions**: -- PI Web API error notification sent -- Workflow fails -- May impact subsequent exports +- PI Web API error notification sent (when applicable) +- Workflow completes +- PostgreSQL contains the prediction with `prediction_confidence` 13 and expected `comments` --- @@ -473,14 +483,15 @@ The `predictions_batch` workflow: - OPC server unavailable or invalid configuration **Expected Behavior**: -- `write_opc_data` raises exception -- Notification sent -- Workflow fails after retries +- `write_opc_data` reports failure without aborting the workflow +- Notification may be sent +- Workflow **completes** (does not fail) +- Prediction row is written to PostgreSQL with OPC error confidence **12** and a comment indicating OPC write issues **Assertions**: -- OPC error notification sent -- Workflow fails -- PostgreSQL export may not execute +- OPC error notification sent (when applicable) +- Workflow completes +- PostgreSQL contains the prediction with `prediction_confidence` 12 and expected `comments` --- @@ -497,7 +508,7 @@ The `predictions_batch` workflow: - `process_pi_web_api_response` detects partial failure - Error confidence set (13) - Notification sent for failed tag -- Workflow completes with error confidence +- Workflow completes with error confidence (single activity attempt; no retry loop) **Assertions**: - One tag written successfully diff --git a/e2e/test_child_workflows_e2e.py b/e2e/test_child_workflows_e2e.py new file mode 100644 index 0000000..52a4b3f --- /dev/null +++ b/e2e/test_child_workflows_e2e.py @@ -0,0 +1,74 @@ +""" +Direct E2E execution of child workflows (smaller surface than PredictionsBatch). +""" + +from decimal import Decimal + +import pytest +from sqlalchemy import text +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from e2e.helpers import make_workflow_id, start_and_await_workflow +from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_format_and_export_prediction_default_path_e2e( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + postgres_engine, +): + """ + Run FormatAndExportPrediction with path_flag set (format_default_prediction path). + """ + client = temporal_test_env.client + model_id = 401 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + + metadata = { + 'metadata': { + 'model_id': model_id, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'subworkflow.format_and_export_prediction', + } + } + input_data = { + 'metadata': metadata, + 'path_flag': 'CONTINUE', + 'data': {'last_timestamp': '2024-01-01 12:00:00+00:00'}, + 'prediction_confidence': 2, + 'timestamp': '2024-01-01 12:00:00+00:00', + 'model_id': model_id, + 'model_name': 'test_model', + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'comment': 'e2e child workflow default path', + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'prediction_store_policy': 'lts:1', + } + + await start_and_await_workflow( + client, + FormatAndExportPrediction.run, + input_data, + make_workflow_id('e2e-format-export-child'), + ) + + with postgres_engine.connect() as conn: + row = conn.execute( + text( + f'SELECT prediction, prediction_confidence, prediction_status, comments ' + f'FROM predictions_schema.predictions WHERE model_id = {model_id}' + ) + ).fetchone() + assert row is not None + assert row[0] == 0 + assert row[1] == Decimal(2) + assert row[2] == 'Bad' + assert row[3] == 'e2e child workflow default path' diff --git a/e2e/test_minio_offload.py b/e2e/test_minio_offload.py new file mode 100644 index 0000000..f8c0122 --- /dev/null +++ b/e2e/test_minio_offload.py @@ -0,0 +1,124 @@ +""" +E2E-style tests for MinIO offload using a real MinIO testcontainer. +""" + +from unittest.mock import patch + +import pytest +from sqlalchemy import text +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from e2e.helpers import insert_sample_data, make_workflow_id, start_and_await_workflow +from laborious.activities.activities import Activities +from laborious.utils.models import minio_dataframe_payload as mdp +from laborious.workflows.predictions_batch import PredictionsBatch + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_load_query_with_minio_offload_writes_object_to_bucket( + postgres_engine, + minio_container, + test_activities_real_minio: Activities, +): + """ + With a tiny offload threshold, query results are uploaded as Parquet to MinIO. + + Uses real MinioRepository against testcontainers MinIO (no MinIO mock). + """ + model_id = 501 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + insert_sample_data(postgres_engine, model_id, [1.0, 2.0]) + + metadata = { + 'metadata': { + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': model_id, + 'workflow_name': 'predictions_batch', + } + } + with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1): + payload = await test_activities_real_minio.load_query_with_minio_offload( + { + **metadata, + 'query': ( + 'SELECT timestamp, variable, value, created_at ' + f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + ), + 'model_name': 'test_model', + 'datetime_columns': ['timestamp', 'created_at'], + } + ) + assert payload.object_key, 'offloaded payload must reference a MinIO object' + assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict' + + df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata']) + assert len(df) >= 1 + + client = minio_container.get_client() + listed = list(client.list_objects('test-bucket', recursive=True)) + names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed] + assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}' + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_predictions_batch_with_minio_offload_path( + temporal_test_env: WorkflowEnvironment, + temporal_worker_real_minio: Worker, + postgres_engine, + test_activities_real_minio: Activities, +): + """ + Full PredictionsBatch run with offload: load step stores Parquet in MinIO; pipeline completes. + """ + model_id = 502 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')) + insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0]) + + input_data = { + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': model_id, + 'query': ( + 'SELECT timestamp, variable, value, created_at ' + f'FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + ), + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'input_filters': {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'mlflow_transform_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'mlflow_predict_filters': {'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}}, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'save_transform': True, + 'prediction_store_policy': 'lts:1', + 'model_config': { + 'retention_minutes': 0, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'sklearn', + }, + 'datetime_columns': ['timestamp', 'created_at'], + } + + with patch.object(mdp, 'OFFLOAD_THRESHOLD_BYTES', 1): + await start_and_await_workflow( + temporal_test_env.client, + PredictionsBatch.run, + input_data, + make_workflow_id('test-batch-minio-offload'), + ) + + with postgres_engine.connect() as conn: + count = conn.execute( + text(f'SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}') + ).scalar() + assert count == 1 diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py index d7e7621..35543f8 100644 --- a/e2e/test_predictions_batch_format_export.py +++ b/e2e/test_predictions_batch_format_export.py @@ -2,18 +2,17 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios. """ -import asyncio -from datetime import datetime +from decimal import Decimal from typing import Any, cast -from unittest.mock import ANY, AsyncMock, patch, call +from unittest.mock import ANY, AsyncMock, call -import pandas as pd import pytest from sientia_do.notifications.models import NotificationLevel from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch @@ -56,70 +55,6 @@ def get_base_input_data(model_id): 'query': base_query.format(model_id=model_id), } -def insert_sample_data(postgres_engine, model_id, values: list): - with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}")) - - values_sql = [] - for i, value in enumerate(values): - values_sql.append(f""" - ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') - """) - - insert_sql = f""" - INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - {', '.join(values_sql)} - """ - conn.execute(text(insert_sql)) - -async def start_and_await_workflow(client, input_data, workflow_id): - handle = await client.start_workflow( - PredictionsBatch.run, - input_data, - id=workflow_id, - task_queue='test-queue', - ) - print("[TEST] ✓ Workflow started") - - print("\n[TEST] 3. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - pytest.fail("Workflow execution timed out after 60 seconds") - -def assert_prediction( - postgres_engine, model_id, prediction: float = 0.5, - prediction_confidence: int = 0, prediction_status: str = 'Good', - comments: str = '', -): - """ - Verify prediction was created with correct values in database - - Args: - postgres_engine: Database engine - model_id: Model ID to check - prediction: Expected prediction value (default 0.5 from mock) - prediction_confidence: Expected confidence value (default 0 for normal predictions) - prediction_status: Expected status (default 'Good') - comments: Expected comments (default empty string) - """ - print("\n[TEST] 4. Verifying prediction was created with correct values...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - prediction_rows = result_query.fetchall() - assert len(prediction_rows) == 1, f"Expected one prediction record, got {len(prediction_rows)}" - - row = prediction_rows[0] - assert row[0] == model_id, f"Expected model_id={model_id}, got {row[0]}" - assert row[1] == prediction, f"Expected prediction={prediction}, got {row[1]}" - assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}" - assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}" - assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" - @pytest.mark.asyncio @pytest.mark.integration @@ -130,35 +65,31 @@ async def test_scenario_3_1_1_default_prediction_export( postgres_engine, ): """ - Scenario 3.1.1: Default Prediction Export - - Description: - Error prediction path creates default prediction. - - Expected Behavior: - - format_default_prediction called instead of format_prediction - - Default prediction created with error metadata - - Exported to PostgreSQL only - - Transformed data NOT processed - - Metrics written - - Assertions: - - format_default_prediction called - - format_prediction NOT called - - format_transformed_data NOT called - - One PostgreSQL export only - - Default values in prediction data - - Comment included + Scenario 3.1.1: Default prediction export (non-None path_flag). + + Triggers input_gate CONTINUE via SPECIFIC_VARIABLES_NULL_VALUES so + PredictionProcess calls FormatAndExportPrediction with path_flag set. + That workflow uses format_default_prediction (not format_prediction) and + skips format_transformed_data / transform Postgres export. + + Optional PI Web API and OPC outputs still run when configured. """ client = temporal_test_env.client model_id = 311 - print("\n[TEST] 1. Inserting test data...") - insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") + with postgres_engine.begin() as conn: + conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}")) + conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) + insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) input_data = get_base_input_data(model_id) + input_data['input_filters'] = { + 'SPECIFIC_VARIABLES_NULL_VALUES': { + 'POLICY': 'CONTINUE', + 'CONFIG': {'variables': ['sensor_1']}, + }, + } input_data['pi_web_api_output_config'] = { 'endpoint': 'test_endpoint', 'prediction_tags': {'tag_1': 'web_id_1'}, @@ -179,10 +110,9 @@ async def test_scenario_3_1_1_default_prediction_export( } } - print("\n[TEST] 2. Starting workflow that should create default prediction...") - workflow_id = f'test-default-prediction-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + wid = make_workflow_id('test-default-prediction') + + await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -190,7 +120,7 @@ async def test_scenario_3_1_1_default_prediction_export( web_ids=['web_id_1'], value={ 'Timestamp': '2024-01-01 12:00:00+0000', - 'Value': 0.5, + 'Value': 0, }, metadata={ 'model_id': 311, @@ -203,7 +133,7 @@ async def test_scenario_3_1_1_default_prediction_export( web_ids=['web_id_2'], value={ 'Timestamp': '2024-01-01 12:00:00+0000', - 'Value': 0, + 'Value': 2, }, metadata={ 'model_id': 311, @@ -219,26 +149,48 @@ async def test_scenario_3_1_1_default_prediction_export( opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) opc_write_data.assert_has_calls( [ - call('addr_1', 0.5, 'float', ANY, - { - 'model_id': 311, - 'model_name': 'test_model', - 'schedule_name': 'test-schedule', - 'workflow_name': 'predictions_batch', - }), - call('addr_2', 0, 'float', ANY, - { - 'model_id': 311, - 'model_name': 'test_model', - 'schedule_name': 'test-schedule', - 'workflow_name': 'predictions_batch', - }), + call( + 'addr_1', + 0, + 'float', + ANY, + { + 'model_id': 311, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'predictions_batch', + }, + ), + call( + 'addr_2', + 2, + 'float', + ANY, + { + 'model_id': 311, + 'model_name': 'test_model', + 'schedule_name': 'test-schedule', + 'workflow_name': 'predictions_batch', + }, + ), ] ) - assert_prediction(postgres_engine, model_id) + with postgres_engine.connect() as conn: + tf_count = conn.execute( + text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}") + ).scalar() + assert tf_count == 0, 'transform export must be skipped when path_flag is set' + + assert_prediction( + postgres_engine, + model_id, + prediction=0, + prediction_confidence=Decimal(2), + prediction_status='Bad', + comments='Input data with bad quality', + ) - print("\n[TEST] ✓ All assertions passed!") @@ -274,9 +226,7 @@ async def test_scenario_3_1_2_export_with_opc_only( model_id = 312 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['opc_output_config'] = { @@ -295,10 +245,9 @@ async def test_scenario_3_1_2_export_with_opc_only( } input_data['pi_web_api_output_config'] = None # No PI Web API config - print("\n[TEST] 2. Starting workflow with OPC only...") - workflow_id = f'test-opc-only-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-only') + ) opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) opc_write_data.assert_has_calls( @@ -324,7 +273,6 @@ async def test_scenario_3_1_2_export_with_opc_only( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -358,9 +306,7 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( model_id = 313 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['pi_web_api_output_config'] = { @@ -370,10 +316,9 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( } input_data['opc_output_config'] = None # No OPC config - print("\n[TEST] 2. Starting workflow with PI Web API only...") - workflow_id = f'test-pi-api-only-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-only') + ) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -412,7 +357,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -445,18 +389,15 @@ async def test_scenario_3_1_4_export_without_optional_outputs( model_id = 314 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['opc_output_config'] = None # No OPC config input_data['pi_web_api_output_config'] = None # No PI Web API config - print("\n[TEST] 2. Starting workflow without optional outputs...") - workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-no-optional-outputs') + ) test_activities.pi_web_api_client.write_value.assert_not_called() opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) @@ -464,7 +405,6 @@ async def test_scenario_3_1_4_export_without_optional_outputs( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -495,11 +435,9 @@ async def test_scenario_3_1_5_export_without_transformed_data( model_id = 315 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) with postgres_engine.begin() as conn: conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) - print("[TEST] ✓ Data inserted successfully") input_data = get_base_input_data(model_id) input_data['save_transform'] = False # Don't save transformed data @@ -523,10 +461,9 @@ async def test_scenario_3_1_5_export_without_transformed_data( } } - print("\n[TEST] 2. Starting workflow without transformed data export...") - workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-no-transform-export') + ) test_activities.pi_web_api_client.write_value.assert_has_calls( [ @@ -589,7 +526,6 @@ async def test_scenario_3_1_5_export_without_transformed_data( assert_prediction(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -599,31 +535,20 @@ async def test_scenario_3_2_1_pi_web_api_write_error( temporal_worker: Worker, test_activities: Activities, postgres_engine, + notification_inserts, ): """ Scenario 3.2.1: PI Web API Write Error - - Description: - PI Web API export fails. - - Expected Behavior: - - write_pi_web_api_data raises exception - - Notification sent - - Workflow fails after retries - - PostgreSQL export may not execute (depends on execution order) - - Assertions: - - PI Web API error notification sent - - Workflow fails - - May impact subsequent exports + + Export failure is handled inside the activity; there is no retry loop. The + workflow completes and PostgreSQL stores prediction_confidence 13 and the + error message in comments. """ client = temporal_test_env.client model_id = 321 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") test_activities.pi_web_api_client.write_value.side_effect = Exception( "PI Web API service unavailable") @@ -649,18 +574,17 @@ async def test_scenario_3_2_1_pi_web_api_write_error( } } - print("\n[TEST] 2. Starting workflow that should fail on PI Web API write...") - workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-error') + ) assert_prediction( postgres_engine, model_id, prediction_confidence=13, comments='PI Web API service unavailable', ) - - print("\n[TEST] ✓ All assertions passed!") + assert notification_inserts.call_count >= 1 + @pytest.mark.asyncio @@ -673,27 +597,15 @@ async def test_scenario_3_2_2_opc_write_error( ): """ Scenario 3.2.2: OPC Write Error - - Description: - OPC server write fails. - - Expected Behavior: - - write_opc_data raises exception - - Notification sent - - Workflow fails after retries - - Assertions: - - OPC error notification sent - - Workflow fails - - PostgreSQL export may not execute + + OPC failure is reported without failing the workflow; there is no retry + loop. PostgreSQL stores prediction_confidence 12 and OPC error comments. """ client = temporal_test_env.client model_id = 322 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) opc_write_data.return_value = (False, { @@ -725,10 +637,9 @@ async def test_scenario_3_2_2_opc_write_error( 'confidence_tags': {'tag_2': 'web_id_2'}, } - print("\n[TEST] 2. Starting workflow that should fail on OPC write...") - workflow_id = f'test-opc-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error') + ) assert_prediction( postgres_engine, model_id, @@ -736,7 +647,6 @@ async def test_scenario_3_2_2_opc_write_error( comments='Some data could not be written to OPC servers', ) - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -749,31 +659,15 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( ): """ Scenario 3.2.3: PI Web API Partial Write Error - - Description: - Two prediction tags attempt to be written to PI Web API, but only one succeeds. - - Expected Behavior: - - write_pi_web_api_data processes response - - process_pi_web_api_response detects partial failure - - Error confidence set (13) - - Notification sent for failed tag - - Workflow completes with error confidence - - Assertions: - - One tag written successfully - - One tag failed - - Error confidence set in prediction - - Error notification sent - - Workflow completes + + Partial PI write: confidence 13, descriptive comments, workflow completes + without an activity retry loop. """ client = temporal_test_env.client model_id = 323 - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") test_activities.pi_web_api_client.write_value = AsyncMock( side_effect=[ @@ -805,10 +699,9 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( } } - print("\n[TEST] 2. Starting workflow with partial PI Web API write error...") - workflow_id = f'test-pi-api-partial-error-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error') + ) assert_prediction( postgres_engine, model_id, @@ -816,4 +709,3 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.", ) - print("\n[TEST] ✓ All assertions passed!") diff --git a/e2e/test_predictions_batch_main_workflow.py b/e2e/test_predictions_batch_main_workflow.py index ddac7ec..3f8d37f 100644 --- a/e2e/test_predictions_batch_main_workflow.py +++ b/e2e/test_predictions_batch_main_workflow.py @@ -3,14 +3,14 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios. """ import asyncio -from datetime import datetime -import pandas as pd -import pytest from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +import pytest + +from e2e.helpers import make_workflow_id, start_and_await_workflow from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch @@ -23,66 +23,20 @@ async def test_scenario_1_1_1_happy_path_complete_success( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.1.1: Happy Path - Complete Success - - Description: - Workflow completes successfully with valid SQL query and all activities succeed. - - Process Flow: - 1. load_custom_query returns DataFrame with sensor data - 2. Workflow prepares prediction input with all configurations - 3. prediction_process child workflow executes: - - get_last_timestamp retrieves last processing timestamp - - input_gate validates data quality (passes) - - request_transform calls MLFlow transform (mocked, returns features) - - mlflow_response_gate validates transform response (passes) - - mlflow_content_gate validates transform content (passes) - - request_predict calls MLFlow predict (mocked, returns predictions) - - mlflow_response_gate validates predict response (passes) - - mlflow_content_gate validates predict content (passes) - 4. format_and_export_prediction child workflow executes: - - format_prediction formats the prediction data - - format_transformed_data formats transformed data (if save_transform=True) - - export_data_to_postgres saves to database - - write_metrics records execution metrics - - Expected Behavior: - - All activities execute successfully without errors - - All gates pass with no quality issues - - Transform and predict operations succeed (mocked) - - Data exported to PostgreSQL predictions table - - Transformed data exported to transformed_data table (if save_transform=True) - - Metrics written successfully - - Assertions: - - Workflow completes without raising exceptions - - Data exists in PostgreSQL predictions table with correct model_id - - Data exists in transformed_data table (if save_transform=True) - - Prediction data has expected structure (jsonb with predictions) - - All required fields are populated (model_id, model_name, timestamp, etc) - """ + """Scenario 1.1.1: Happy path with SQL load, MLflow mocks, Postgres predictions and transforms.""" client = temporal_test_env.client - print("\n[TEST] 1. Inserting test data into PostgreSQL...") - # Insert test data directly into PostgreSQL - # The load_custom_query activity will fetch this data with a real SQL query with postgres_engine.begin() as conn: - # Clear any existing data for this model_id - conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 123")) - - # Insert sensor data that the workflow will query + conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 123')) insert_sql = """ INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES + VALUES (123, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'), (123, 'sensor_2', 78.2, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00'), (123, 'sensor_3', 120.8, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') """ conn.execute(text(insert_sql)) - print("[TEST] ✓ Data inserted successfully") - # Prepare input data for PredictionsBatch workflow input_data = { 'metadata': { 'metadata': { @@ -121,77 +75,44 @@ async def test_scenario_1_1_1_happy_path_complete_success( 'datetime_columns': ['timestamp', 'created_at'], } - # Start workflow - print("\n[TEST] 2. Starting workflow...") - workflow_id = f'test-predictions-batch-{datetime.now().timestamp()}' - print(f"[TEST] Workflow ID: {workflow_id}") - - handle = await client.start_workflow( + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, - id=workflow_id, - task_queue='test-queue', + make_workflow_id('test-predictions-batch'), ) - print("[TEST] ✓ Workflow started") - # Wait for workflow completion with timeout - print("\n[TEST] 3. Waiting for workflow completion (timeout: 60s)...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) # 60 seconds timeout - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - print("[TEST] ✗ Workflow TIMEOUT after 60 seconds!") - pytest.fail("Workflow execution timed out after 60 seconds") - - # Verify data was stored in PostgreSQL - use single connection schema_name = 'predictions_schema' - predictions_table = 'predictions' - transformed_table = 'transformed_data' - full_predictions_table = f"{schema_name}.{predictions_table}" - full_transformed_table = f"{schema_name}.{transformed_table}" - - # Use a single connection for all verification queries - print("\n[TEST] 4. Verifying results in PostgreSQL...") with postgres_engine.connect() as conn: - # Verify prediction data result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments FROM {full_predictions_table} WHERE model_id = 123") + text( + f'SELECT model_id, prediction, prediction_confidence, response_time, prediction_status, comments ' + f'FROM {schema_name}.predictions WHERE model_id = 123' + ) ) prediction_rows = result_query.fetchall() - - print(f"[TEST] Found {len(prediction_rows)} prediction record(s)") - assert len(prediction_rows) == 1, "Expected one prediction record" - - # Verify first row has expected structure + assert len(prediction_rows) == 1 row = prediction_rows[0] - print(f"[TEST] Prediction: {row}") - assert row[0] == 123, f"Expected model_id=123, got {row[0]}" - assert row[1] == 0.5, f"Expected prediction=0.5, got {row[1]}" - assert row[2] == 0, f"Expected prediction_confidence=0.9, got {row[2]}" - assert row[3] is not None, f"Expected response_time=0.1, got {row[3]}" - assert row[4] == 'Good', f"Expected prediction_status='Good', got {row[4]}" - assert row[5] == '', f"Expected comments='', got {row[5]}" - print("[TEST] ✓ Prediction data verified") + assert row[0] == 123 + assert row[1] == 0.5 + assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}' + assert row[3] is not None + assert row[4] == 'Good' + assert row[5] == '' - # Verify transformed data result_query = conn.execute( - text(f"SELECT model_id, variable, value FROM {full_transformed_table} WHERE model_id = 123") + text( + f'SELECT model_id, variable, value FROM {schema_name}.transformed_data WHERE model_id = 123' + ) ) transformed_rows = result_query.fetchall() - print(f"[TEST] Found {len(transformed_rows)} transformed data record(s)") - assert len(transformed_rows) == 2, "Expected two transformed data records" - row_1 = transformed_rows[0] - print(f"[TEST] Transformed data: {row_1}") - assert row_1[0] == 123, f"Expected model_id=123, got {row_1[0]}" - assert row_1[1] == 'feature_1', f"Expected variable='sensor_1', got {row_1[1]}" - assert float(row_1[2]) == 0.234, f"Expected value=0.234, got {row_1[2]}" - row_2 = transformed_rows[1] - print(f"[TEST] Transformed data: {row_2}") - assert row_2[0] == 123, f"Expected model_id=123, got {row_2[0]}" - assert row_2[1] == 'feature_2', f"Expected variable='sensor_2', got {row_2[1]}" - assert float(row_2[2]) == 0.783, f"Expected value=0.783, got {row_2[2]}" - - print("\n[TEST] ✓ All assertions passed!") + assert len(transformed_rows) == 2 + assert transformed_rows[0][0] == 123 + assert transformed_rows[0][1] == 'feature_1' + assert float(transformed_rows[0][2]) == 0.234 + assert transformed_rows[1][0] == 123 + assert transformed_rows[1][1] == 'feature_2' + assert float(transformed_rows[1][2]) == 0.783 @pytest.mark.asyncio @@ -202,23 +123,7 @@ async def test_scenario_1_2_1_sql_query_execution_error( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.1: SQL Query Execution Error - - Description: - SQL query fails due to syntax error or connection issue. - - Expected Behavior: - - load_custom_query raises exception (caught by Temporal retry policy) - - Notification sent with SQL error details - - After retries, activity may return empty data or workflow may fail - - If empty data returned, workflow completes with early exit via input gate - - Assertions: - - Error notification sent - - Workflow completes (either fails or exits early) - - No data in predictions table - """ + """Invalid SQL: workflow may complete with early exit; no prediction rows.""" client = temporal_test_env.client input_data = { @@ -233,7 +138,7 @@ async def test_scenario_1_2_1_sql_query_execution_error( 'schedule_name': 'test-schedule', 'model_name': 'test_model', 'model_id': 128, - 'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =', # Invalid SQL + 'query': 'SELECT * FROM nonexistent_table WHERE invalid_syntax =', 'schema': 'predictions_schema', 'table_name': 'predictions', 'transform_table_name': 'transformed_data', @@ -258,34 +163,18 @@ async def test_scenario_1_2_1_sql_query_execution_error( }, } - print("\n[TEST] 1. Starting workflow with invalid SQL query...") - workflow_id = f'test-sql-error-{datetime.now().timestamp()}' - - handle = await client.start_workflow( + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, - id=workflow_id, - task_queue='test-queue', + make_workflow_id('test-sql-error'), ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 2. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed (may have exited early due to empty data)") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - # Verify no predictions were created (regardless of whether workflow failed or exited early) - print("\n[TEST] 3. Verifying no predictions were created...") with postgres_engine.connect() as conn: - result_query = conn.execute( - text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128") - ) - count = result_query.scalar() - assert count == 0, f"Expected no predictions, but found {count} records" - - print("\n[TEST] ✓ All assertions passed!") + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128') + ).scalar() + assert count == 0 @pytest.mark.asyncio @@ -296,24 +185,9 @@ async def test_scenario_1_2_2_missing_required_parameters( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.2: Missing Required Parameters - - Description: - Essential parameters missing from input. - - Expected Behavior: - - Workflow or activity raises KeyError or validation error - - Workflow fails immediately - - Assertions: - - Workflow fails with parameter error - - Error notification sent - - No child workflow called - """ + """Missing query: workflow does not produce predictions and is terminated explicitly.""" client = temporal_test_env.client - # Missing 'query' parameter input_data = { 'metadata': { 'metadata': { @@ -326,31 +200,28 @@ async def test_scenario_1_2_2_missing_required_parameters( 'schedule_name': 'test-schedule', 'model_name': 'test_model', 'model_id': 129, - # 'query' is missing 'schema': 'predictions_schema', 'table_name': 'predictions', 'transform_table_name': 'transformed_data', } - print("\n[TEST] 1. Starting workflow with missing required parameter...") - workflow_id = f'test-missing-param-{datetime.now().timestamp()}' - handle = await client.start_workflow( PredictionsBatch.run, input_data, - id=workflow_id, + id=make_workflow_id('test-missing-param'), task_queue='test-queue', ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 2. Waiting for workflow to fail...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - pytest.fail("Expected workflow to fail, but it completed successfully") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - print("\n[TEST] ✓ All assertions passed!") + # Let Temporal process a few workflow tasks; for this case, result() can hang. + await asyncio.sleep(2.0) + + with postgres_engine.connect() as conn: + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129') + ).scalar() + assert count == 0 + + await handle.terminate('expected failure path in e2e test (missing required parameters)') @pytest.mark.asyncio @@ -361,34 +232,19 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( test_activities: Activities, postgres_engine, ): - """ - Scenario 1.2.3: Invalid Datetime Column Specification - - Description: - Datetime column specified doesn't exist in query results. - - Expected Behavior: - - load_custom_query may raise KeyError or warning - - Depending on implementation, workflow may fail or continue - - Error notification sent - - Assertions: - - Error raised or warning logged - - Workflow behavior depends on error handling policy - """ + """Invalid datetime column: no predictions persisted; workflow terminated after validation.""" client = temporal_test_env.client - print("\n[TEST] 1. Inserting test data...") with postgres_engine.begin() as conn: - conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 130")) - - insert_sql = """ + conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130')) + conn.execute( + text( + """ INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') + VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') """ - conn.execute(text(insert_sql)) - print("[TEST] ✓ Data inserted successfully") + ) + ) input_data = { 'metadata': { @@ -425,26 +281,23 @@ async def test_scenario_1_2_3_invalid_datetime_column_specification( 'transform_flavor': 'sklearn', 'predict_flavor': 'sklearn', }, - 'datetime_columns': ['nonexistent_column'], # Column doesn't exist in query result + 'datetime_columns': ['nonexistent_column'], } - print("\n[TEST] 2. Starting workflow with invalid datetime column...") - workflow_id = f'test-invalid-datetime-col-{datetime.now().timestamp()}' - handle = await client.start_workflow( PredictionsBatch.run, input_data, - id=workflow_id, + id=make_workflow_id('test-invalid-datetime-col'), task_queue='test-queue', ) - print("[TEST] ✓ Workflow started") - print("\n[TEST] 3. Waiting for workflow completion or failure...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - # Workflow may complete or fail depending on error handling - print("[TEST] ✓ Workflow completed (may have handled error gracefully)") - except Exception as e: - print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}") - - print("\n[TEST] ✓ Test completed!") + # Let Temporal process and surface the failure path internally. + await asyncio.sleep(2.0) + + with postgres_engine.connect() as conn: + count = conn.execute( + text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130') + ).scalar() + assert count == 0 + + await handle.terminate('expected failure path in e2e test (invalid datetime column)') diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index 98f9429..af097cd 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/e2e/test_predictions_batch_prediction_process.py @@ -2,56 +2,62 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios. """ -import asyncio -from datetime import datetime from decimal import Decimal -from typing import Any from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd import pytest -from pytz import timezone from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +from e2e.helpers import ( + assert_continue, + assert_repeat, + assert_stop, + insert_sample_data, + make_workflow_id, + start_and_await_workflow, +) from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch base_input_data = { - 'schedule_name': 'test-schedule', - 'model_name': 'test_model', - 'model_id': 201, - 'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201', - 'schema': 'predictions_schema', - 'table_name': 'predictions', - 'transform_table_name': 'transformed_data', - 'input_filters': { - 'SPECIFIC_VARIABLES_NULL_VALUES': { - 'POLICY': 'CONTINUE', # Continue despite issues, not STOP - 'CONFIG': {'variables': ['sensor_1']}, - }, + 'schedule_name': 'test-schedule', + 'model_name': 'test_model', + 'model_id': 201, + 'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201', + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'input_filters': { + 'SPECIFIC_VARIABLES_NULL_VALUES': { + 'POLICY': 'CONTINUE', + 'CONFIG': {'variables': ['sensor_1']}, }, - 'mlflow_transform_filters': { - 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, - }, - 'mlflow_predict_filters': { - 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, - }, - 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], - 'opc_output_config': {}, - 'pi_web_api_output_config': {}, - 'save_transform': True, - 'prediction_store_policy': 'lts:1', - 'model_config': { - 'retention_minutes': 0, - 'transform_flavor': 'sklearn', - 'predict_flavor': 'sklearn', - }, - 'datetime_columns': ['timestamp', 'created_at'], - } + }, + 'mlflow_transform_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'mlflow_predict_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'save_transform': True, + 'prediction_store_policy': 'lts:1', + 'model_config': { + 'retention_minutes': 0, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'sklearn', + }, + 'datetime_columns': ['timestamp', 'created_at'], +} + +base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}' -base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}" def get_base_input_data(model_id): return { @@ -60,97 +66,38 @@ def get_base_input_data(model_id): 'query': base_query.format(model_id=model_id), } -def insert_sample_data(postgres_engine, model_id, values: list[Any]): - with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}")) - - # Insert data with some null values (quality issue) - - values_sql = [] - for i, value in enumerate(values): - values_sql.append(f""" - ({model_id}, 'sensor_{i+1}', {value}, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00') - """) - - insert_sql = f""" - INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at) - VALUES - {', '.join(values_sql)} - """ - conn.execute(text(insert_sql)) - def insert_sample_prediction(postgres_engine, model_id): with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}")) - - # Insert data with some null values (quality issue) - + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) insert_sql = f""" INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time) VALUES ({model_id}, '2024-01-01 12:00:00+00:00', 10, 0, 'Good', '', 0.1) """ conn.execute(text(insert_sql)) - - return (model_id, Decimal(10), Decimal(0), 'Good') - -async def start_and_await_workflow(client, input_data, workflow_id): - handle = await client.start_workflow( - PredictionsBatch.run, - input_data, - id=workflow_id, - task_queue='test-queue', - ) - print("[TEST] ✓ Workflow started") - - print("\n[TEST] 3. Waiting for workflow completion...") - try: - await asyncio.wait_for(handle.result(), timeout=60.0) - print("[TEST] ✓ Workflow completed successfully") - except asyncio.TimeoutError: - pytest.fail("Workflow execution timed out after 60 seconds") - -def assert_continue( - postgres_engine, model_id, prediction_confidence: Decimal = Decimal(2), - comments: str = 'Input data with bad quality', -): - print("\n[TEST] 4. Verifying prediction was created despite warnings...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT model_id, prediction, prediction_confidence, prediction_status, comments FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - prediction_rows = result_query.fetchall() - assert len(prediction_rows) == 1, "Expected one prediction record despite warnings" - - # Assert prediction value is 0 and other fields - row = prediction_rows[0] - assert row[1] == 0, f"Expected prediction=0, got {row[1]}" - assert row[2] == prediction_confidence, f"Expected prediction_confidence={prediction_confidence}, got {row[2]}" - assert row[3] == 'Bad', f"Expected prediction_status='Bad', got {row[3]}" - assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" + return (model_id, Decimal(10), Decimal(0), 'Good') -def assert_stop(postgres_engine, model_id): - print("\n[TEST] 4. Verifying no predictions were created...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}") - ) - count = result_query.scalar() - assert count == 0, f"Expected no predictions, but found {count} records" +@pytest.fixture +def bad_data_model(patch_mlflow): + model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model'))) + patch_mlflow.sklearn.load_model = MagicMock(return_value=model) + return model -def assert_repeat(postgres_engine, model_id, last_prediction: list): - print("\n[TEST] 4. Verifying prediction was repeated...") - with postgres_engine.connect() as conn: - result_query = conn.execute( - text(f'SELECT model_id, prediction, prediction_confidence, prediction_status FROM predictions_schema.predictions WHERE model_id = {model_id}') - ) - prediction_rows = result_query.fetchall() - print(prediction_rows) - assert len(prediction_rows) == 2, "Expected two prediction records" - assert prediction_rows[0] == last_prediction, f"Expected first prediction to be the same as the last prediction, got {prediction_rows[0]}, expected {last_prediction}" - assert prediction_rows[1] == last_prediction, f"Expected second prediction to be the same as the last prediction, got {prediction_rows[1]}, expected {last_prediction}" + +@pytest.fixture +def bad_predict_model(patch_mlflow, mock_mlflow_models): + model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model'))) + + def mock_sklearn_load_model(model_uri): + if 'data_model' in model_uri or 'transform' in model_uri.lower(): + return mock_mlflow_models['transform_model'] + return model + + patch_mlflow.sklearn = MagicMock() + patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model) + return model @pytest.mark.asyncio @@ -160,45 +107,19 @@ async def test_scenario_2_1_1_input_gate_triggers_continue( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.1: Input Gate Triggers CONTINUE - - Description: - Input gate determines data should use previous prediction. - - Expected Behavior: - - input_gate returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with input data directly - - MLFlow transform and predict skipped - - Data exported as-is - - Assertions: - - input_gate called - - MLFlow operations NOT called - - Export workflow called with original data - - Workflow completes - """ + """Input gate CONTINUE: export default prediction; MLflow transform/predict not used.""" client = temporal_test_env.client - model_id = 211 - - print("\n[TEST] 1. Inserting test data...") - insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) - - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) - - - print("\n[TEST] 2. Starting workflow with CONTINUE policy...") - workflow_id = f'test-continue-policy-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy') + ) assert_continue(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") + mock_mlflow_models['transform_model'].predict.assert_not_called() + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -208,45 +129,19 @@ async def test_scenario_2_1_2_input_gate_triggers_stop( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.2: Input Gate Triggers STOP - - Description: - Input data quality gate fails with STOP policy. - - Expected Behavior: - - input_gate returns path_flag='STOP' - - path_flag_handler detects STOP - - Workflow returns early without calling MLFlow - - No prediction exported - - Assertions: - - input_gate called - - path_flag_handler returns True (early exit) - - MLFlow transform NOT called - - Export workflow NOT called - - Workflow completes without error - """ + """Input gate STOP: no export, no MLflow.""" client = temporal_test_env.client - model_id = 212 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should stop at input gate...") - workflow_id = f'test-input-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) - + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop') + ) assert_stop(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") + mock_mlflow_models['transform_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -256,57 +151,42 @@ async def test_scenario_2_1_3_input_gate_triggers_repeat( temporal_worker: Worker, test_activities: Activities, postgres_engine, + mock_mlflow_models, ): - """ - Scenario 2.1.3: Input Gate Triggers REPEAT - - Description: - Input gate determines data should repeat last prediction. - - Expected Behavior: - - input_gate returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - MLFlow transform and predict skipped - - Last prediction repeated and exported - - Assertions: - - input_gate called - - MLFlow operations NOT called - - repeat_last_prediction activity called - - Workflow completes - """ + """Input gate REPEAT with existing history.""" client = temporal_test_env.client - model_id = 213 - - print("\n[TEST] 1. Inserting test data and previous prediction...") insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data and previous prediction inserted") - input_data = get_base_input_data(model_id) input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT...") - workflow_id = f'test-input-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) - assert_repeat(postgres_engine, model_id, data) - - print("\n[TEST] ✓ All assertions passed!") - -@pytest.fixture -def bad_data_model(patch_mlflow): - model = MagicMock( - predict=MagicMock( - side_effect=Exception("Bad data model") - ) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat') ) + assert_repeat(postgres_engine, model_id, data) + mock_mlflow_models['transform_model'].predict.assert_not_called() - patch_mlflow.sklearn.load_model = MagicMock(return_value=model) - - return model + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, +): + """REPEAT when no prior row in predictions: repeat_last_prediction runs; still no new duplicate export path.""" + client = temporal_test_env.client + model_id = 214 + insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + input_data = get_base_input_data(model_id) + input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-input-repeat-no-history') + ) + assert_stop(postgres_engine, model_id) @pytest.mark.asyncio @@ -318,48 +198,20 @@ async def test_scenario_2_2_1_transform_gate_triggers_continue( postgres_engine, bad_data_model, ): - """ - Scenario 2.2.1: Transform Gate Triggers CONTINUE - - Description: - Transform response gate determines data should continue despite issues. - - Expected Behavior: - - request_transform succeeds - - mlflow_response_gate for transform returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with transform data - - MLFlow predict skipped - - Transform data exported as-is - - Assertions: - - Transform completed - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - Export workflow called with transform data - - Workflow completes - """ client = temporal_test_env.client - model_id = 221 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - input_data = get_base_input_data(model_id) input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE' - - print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...") - workflow_id = f'test-transform-continue-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue') + ) assert_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), comments='Unknown MLFlow API error', ) - - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -370,44 +222,18 @@ async def test_scenario_2_2_2_transform_gate_triggers_stop( test_activities: Activities, postgres_engine, bad_data_model, + mock_mlflow_models, ): - """ - Scenario 2.2.2: Transform Gate Triggers STOP - - Description: - Transform response validation fails with STOP policy. - - Expected Behavior: - - request_transform succeeds but response invalid - - mlflow_response_gate for transform returns path_flag='STOP' - - Workflow exits without calling predict or export - - Assertions: - - Transform completed but validation failed - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - Export workflow NOT called - - Workflow completes without error - """ client = temporal_test_env.client - model_id = 222 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...") - workflow_id = f'test-transform-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop') + ) assert_stop(postgres_engine, model_id) - - print("\n[TEST] ✓ All assertions passed!") - + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -419,67 +245,50 @@ async def test_scenario_2_2_3_transform_gate_triggers_repeat( postgres_engine, bad_data_model, ): - """ - Scenario 2.2.3: Transform Gate Triggers REPEAT - - Description: - Transform response gate determines data should repeat last prediction. - - Expected Behavior: - - request_transform succeeds but response has issues - - mlflow_response_gate for transform returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - MLFlow predict skipped - - Last prediction repeated and exported - - Assertions: - - Transform completed but validation triggered REPEAT - - mlflow_response_gate called for transform - - MLFlow predict NOT called - - repeat_last_prediction activity called - - Workflow completes - """ client = temporal_test_env.client - model_id = 223 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT' - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...") - workflow_id = f'test-transform-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat') + ) assert_repeat(postgres_engine, model_id, data) - print("\n[TEST] ✓ All assertions passed!") - -@pytest.fixture -def bad_predict_model( - patch_mlflow, - mock_mlflow_models +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_2_4_transform_content_gate_nan_values_stop( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + mock_mlflow_models, ): - model = MagicMock( - predict=MagicMock( - side_effect=Exception("Bad predict model") - ) + """mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter).""" + client = temporal_test_env.client + model_id = 224 + + def all_nan_transform(data): + num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1 + result = pd.DataFrame({'feature_1': [np.nan] * num_rows, 'feature_2': [np.nan] * num_rows}) + result.index = data.index + return result + + mock_mlflow_models['transform_model'].predict = MagicMock(side_effect=all_nan_transform) + + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = get_base_input_data(model_id) + input_data['mlflow_transform_filters'] = { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + 'NAN_VALUES': {'POLICY': 'STOP', 'CONFIG': {}}, + } + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-content-stop') ) - - def mock_sklearn_load_model(model_uri): - if 'data_model' in model_uri or 'transform' in model_uri.lower(): - return mock_mlflow_models['transform_model'] - return model - patch_mlflow.sklearn = MagicMock() - patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model) - - return model + assert_stop(postgres_engine, model_id) + mock_mlflow_models['predict_model'].predict.assert_not_called() @pytest.mark.asyncio @@ -491,46 +300,20 @@ async def test_scenario_2_3_1_predict_gate_triggers_continue( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.1: Predict Gate Triggers CONTINUE - - Description: - Predict response gate determines data should continue despite issues. - - Expected Behavior: - - request_predict succeeds - - mlflow_response_gate for predict returns path_flag='CONTINUE' - - path_flag_handler calls export workflow with predict data - - Prediction exported despite quality issues - - Assertions: - - Transform and predict completed - - mlflow_response_gate called for predict - - Export workflow called with predict data - - Workflow completes - """ client = temporal_test_env.client - model_id = 231 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) - input_data = get_base_input_data(model_id) input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE' - - print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...") - workflow_id = f'test-predict-continue-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue') + ) assert_continue( postgres_engine=postgres_engine, model_id=model_id, prediction_confidence=Decimal(10), comments='Unknown MLFlow API error', ) - - print("\n[TEST] ✓ All assertions passed!") @pytest.mark.asyncio @@ -542,42 +325,16 @@ async def test_scenario_2_3_2_predict_gate_triggers_stop( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.2: Predict Gate Triggers STOP - - Description: - Prediction validation fails with STOP policy. - - Expected Behavior: - - request_predict succeeds but response invalid - - mlflow_response_gate for predict returns path_flag='STOP' - - Workflow exits without export - - Assertions: - - Transform completed - - Predict completed but validation failed - - Export workflow NOT called - - Workflow completes without error - """ client = temporal_test_env.client - model_id = 232 - - print("\n[TEST] 1. Inserting test data...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) - print("[TEST] ✓ Data inserted successfully") - input_data = get_base_input_data(model_id) input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP' - - print("\n[TEST] 2. Starting workflow that should stop at predict gate...") - workflow_id = f'test-predict-stop-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop') + ) assert_stop(postgres_engine, model_id) - print("\n[TEST] ✓ All assertions passed!") - @pytest.mark.asyncio @pytest.mark.integration @@ -588,43 +345,35 @@ async def test_scenario_2_3_3_predict_gate_triggers_repeat( postgres_engine, bad_predict_model, ): - """ - Scenario 2.3.3: Predict Gate Triggers REPEAT - - Description: - Predict response gate determines data should repeat last prediction. - - Expected Behavior: - - request_predict succeeds but response has issues - - mlflow_response_gate for predict returns path_flag='REPEAT' - - path_flag_handler calls repeat_last_prediction activity - - Last prediction repeated and exported - - Assertions: - - Transform and predict completed but validation triggered REPEAT - - mlflow_response_gate called for predict - - repeat_last_prediction activity called - - Export workflow NOT called with current prediction - - Workflow completes - """ client = temporal_test_env.client - model_id = 233 - - print("\n[TEST] 1. Inserting test data and previous prediction...") insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) data = insert_sample_prediction(postgres_engine, model_id) - - print("[TEST] ✓ Data and previous prediction inserted") - input_data = get_base_input_data(model_id) input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT' - input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] # REPEAT first - - print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...") - workflow_id = f'test-predict-repeat-{datetime.now().timestamp()}' - - await start_and_await_workflow(client, input_data, workflow_id) + input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE'] + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat') + ) assert_repeat(postgres_engine, model_id, data) - print("\n[TEST] ✓ All assertions passed!") + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_scenario_2_4_1_input_empty_data_stop( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, +): + """EMPTY_DATA filter with STOP when query returns no rows (offload payload empty).""" + client = temporal_test_env.client + model_id = 241 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + input_data = get_base_input_data(model_id) + input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}} + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-empty-data-stop') + ) + assert_stop(postgres_engine, model_id) diff --git a/laborious/activities/api.py b/laborious/activities/api.py index c568e98..d91b0e3 100644 --- a/laborious/activities/api.py +++ b/laborious/activities/api.py @@ -248,8 +248,13 @@ class API(SientiaMonitoring): metadata=metadata, ) - data['prediction_confidence'] = confidence - data['comments'] = message + # 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() diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 58b60c9..6a6207a 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -1,5 +1,3 @@ -import json - from temporalio import activity, workflow from laborious.utils.repository.minio_manager import MinioManager @@ -195,14 +193,9 @@ class Storage(Postgres, MinioManager): ) self.error(trace, metadata) else: - await self.send_notification_async( - metadata=metadata, - notification_id='CLEANUP_MINIO_OBJECTS_EXPIRED', - message='MinIO objects cleaned up successfully', - block='cleanup_minio_objects_expired', - level=NotificationLevel.INFO, - attachment_content=json.dumps(report), - ) + # 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 diff --git a/requirements-dev.txt b/requirements-dev.txt index 16a8492..c4f4ef6 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -13,7 +13,7 @@ types-requests>=2.31.0 # Type stubs for requests pytest>=7.4.0 # Testing framework pytest-cov>=4.1.0 # Coverage plugin for pytest pytest-asyncio>=0.21.0 # Async test support (already in main requirements) -testcontainers[postgres] # PostgreSQL containers for E2E tests +testcontainers[postgres,minio] # PostgreSQL and MinIO containers for E2E tests # Development Tools ipython>=8.12.0 # Enhanced Python shell From 39f3a08e930f6e6f4e631cf294e84beb75e3f615 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 10:18:03 -0300 Subject: [PATCH 24/51] SIENTIAPDE-1712 Update README.md to add newline at the end of the note section for consistency --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index c9c37d6..b63be57 100644 --- a/README.md +++ b/README.md @@ -1195,4 +1195,4 @@ For support and questions: --- -**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. \ No newline at end of file +**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments. From 9fa8a6bf50cfa30db5f9ffbc85820b033913ad76 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 10:29:26 -0300 Subject: [PATCH 25/51] SIENTIAPDE-1712 Update pyproject.toml to exclude specific test file from linting --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index daaeabe..5123123 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ exclude = [ "*.pyc", ".pytest_cache", "htmlcov", + "tests/laborious/workflows/subworkflows/test_prediction_process.py", ] [tool.ruff.lint] From d3f8bf7196d0e6b146cf86f7257944452a854f43 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 10:34:46 -0300 Subject: [PATCH 26/51] SIENTIAPDE-1712 Update pyproject.toml to remove specific test file from exclusion list and bump sientia-dataops-library dependency version to 1.10.4 in requirements-light.txt --- pyproject.toml | 1 - requirements-light.txt | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 5123123..daaeabe 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,6 @@ exclude = [ "*.pyc", ".pytest_cache", "htmlcov", - "tests/laborious/workflows/subworkflows/test_prediction_process.py", ] [tool.ruff.lint] diff --git a/requirements-light.txt b/requirements-light.txt index 55788d0..91a9c33 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -3,7 +3,7 @@ psycopg2-binary sqlalchemy asyncua redis -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.8.0 +git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4 prometheus-client botocore boto3 From a3dfb1ffc6b9cad8c6c7cb0514b10fd3be36a60b Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 10:37:15 -0300 Subject: [PATCH 27/51] SIENTIAPDE-1712 Update pyproject.toml to add specific test file to exclusion list for linting --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index daaeabe..5123123 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,6 +23,7 @@ exclude = [ "*.pyc", ".pytest_cache", "htmlcov", + "tests/laborious/workflows/subworkflows/test_prediction_process.py", ] [tool.ruff.lint] From 1f0441c7d76b17e1f5f886794ee1468327d53e90 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 24 Mar 2026 14:45:22 -0300 Subject: [PATCH 28/51] SIENTIAPDE-1712 Update requirements.txt to add scikit-learn version 1.5.2 and ensure pycurl is included --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index f0eedf8..7cb8e5b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,4 +13,5 @@ pyarrow kaleido hyperopt shap -pycurl \ No newline at end of file +pycurl +scikit-learn==1.5.2 \ No newline at end of file From 367b00f00a0551ffa8759b00f4bc539890f3e567 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 25 Mar 2026 09:41:12 -0300 Subject: [PATCH 29/51] SIENTIAPDE-1712 Enhance prediction export functionality by adding 'on_conflict' and 'unique_columns' parameters to the payload initialization in FormatAndExportPrediction class. --- .../workflows/sub_workflows/format_and_export_prediction.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index c9127c5..a9cda2a 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -198,6 +198,8 @@ class FormatAndExportPrediction: 'table_name': input_data['table_name'], 'data': prediction, 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, + 'on_conflict': input_data.get('on_conflict', 'replace'), + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=180), From 1f5a14be3fdafd3dea60be238eff2f0a80ee3147 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 25 Mar 2026 11:18:21 -0300 Subject: [PATCH 30/51] SIENTIAPDE-1712 Add debug logging for non-DataFrame prediction data in MLFlowRepository --- laborious/utils/repository/model_repository.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index b5f9e9b..fc965dd 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1283,6 +1283,10 @@ class MLFlowRepository(SientiaMonitoring): 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) From 5b4903fffa669eb036c3af17a0553fae90592907 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 25 Mar 2026 12:18:27 -0300 Subject: [PATCH 31/51] SIENTIAPDE-1712 Update Gates and FormatAndExportPrediction classes to use 'last_timestamp' for improved data handling - Modified the Gates class to utilize 'last_timestamp' when only one row is present, ensuring accurate timestamp assignment. - Updated the FormatAndExportPrediction class to replace 'timestamp' with 'last_timestamp' in the output data structure. --- laborious/activities/gates.py | 11 +++++++++-- .../sub_workflows/format_and_export_prediction.py | 2 +- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 1d561c5..0e8dc82 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -514,6 +514,8 @@ class Gates(MinioManager): 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 @@ -521,7 +523,7 @@ class Gates(MinioManager): Args: input_data (dict): Input data containing: - data (dict[str, Any]): Raw prediction data to format - - timestamp (str): Default timestamp if data lacks timestamp column + - last_timestamp (str): Last 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' @@ -530,6 +532,7 @@ class Gates(MinioManager): dict: Formatted prediction data ready for storage and export """ metadata = input_data['metadata'] + last_timestamp = input_data['last_timestamp'] prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) @@ -564,7 +567,11 @@ class Gates(MinioManager): self.error(f'Invalid policy type: {policy_type}, using default policy', metadata) raise ValueError(f'Invalid policy type: {policy_type}') - data = data.head(int(policy_value)) + int_policy_value = int(policy_value) + + data = data.head(int_policy_value) + if int_policy_value == 1: + data['timestamp'] = last_timestamp data['model_id'] = input_data['model_id'] data['prediction_confidence'] = input_data['prediction_confidence'] diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index a9cda2a..1cd2be1 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -101,7 +101,7 @@ class FormatAndExportPrediction: { **metadata, 'data': data, - 'timestamp': input_data['timestamp'], + 'last_timestamp': input_data['last_timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, 'prediction_store_policy': input_data['prediction_store_policy'], From 830dfe464ef08284b6e7e1abaab642d8431c9682 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 25 Mar 2026 12:45:44 -0300 Subject: [PATCH 32/51] SIENTIAPDE-1712 Update Gates and FormatAndExportPrediction classes to standardize timestamp usage - Changed 'last_timestamp' to 'timestamp' in the Gates class for consistency in input data handling. - Updated the FormatAndExportPrediction class to reflect the same change in the output data structure. --- laborious/activities/gates.py | 4 ++-- .../workflows/sub_workflows/format_and_export_prediction.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 0e8dc82..12fafdf 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -523,7 +523,7 @@ class Gates(MinioManager): Args: input_data (dict): Input data containing: - data (dict[str, Any]): Raw prediction data to format - - last_timestamp (str): Last timestamp of the data + - 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' @@ -532,7 +532,7 @@ class Gates(MinioManager): dict: Formatted prediction data ready for storage and export """ metadata = input_data['metadata'] - last_timestamp = input_data['last_timestamp'] + last_timestamp = input_data['timestamp'] prediction_store_policy = input_data['prediction_store_policy'] self.info('Formatting prediction...', metadata) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 1cd2be1..a9cda2a 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -101,7 +101,7 @@ class FormatAndExportPrediction: { **metadata, 'data': data, - 'last_timestamp': input_data['last_timestamp'], + 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, 'prediction_store_policy': input_data['prediction_store_policy'], From 0bba36e0880165c94fb8aef34e20ad320eab4311 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 26 Mar 2026 14:16:13 -0300 Subject: [PATCH 33/51] SIENTIAPDE-1712 Update values.yaml to enhance PredictionsBatch worker tuning parameters - Renamed SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to SIE_MINIO_OFFLOAD_THRESHOLD_MEGABYTES for consistency. - Added multiple configuration options for PredictionsBatch to optimize workflow and activity concurrency, including max concurrent tasks, activities, local activities, cached workflows, and poller behaviors. --- values.yaml | 35 ++++++++++++++++++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/values.yaml b/values.yaml index e7c32f9..51e96fb 100644 --- a/values.yaml +++ b/values.yaml @@ -231,9 +231,42 @@ env: value: "sientia" - name: MINIO_RETENTION_HOURS value: "24" - - name: SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES + - name: SIE_MINIO_OFFLOAD_THRESHOLD_MEGABYTES value: "1.5" + # Temporal worker tuning for PredictionsBatch. + # IMPORTANT: prefix must be PREDICTIONSBATCH_ (from class name PredictionsBatch). + # Keep workflow-task concurrency moderate to reduce task completion races under load. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_WORKFLOW_TASKS + value: "20" + # Allow higher activity parallelism because most activities are I/O-bound, but keep headroom. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_ACTIVITIES + value: "60" + # Keep local activities controlled so they do not monopolize the event loop. + - name: PREDICTIONSBATCH_MAX_CONCURRENT_LOCAL_ACTIVITIES + value: "20" + # Cache enough workflows for reuse without excessive memory growth. + - name: PREDICTIONSBATCH_MAX_CACHED_WORKFLOWS + value: "200" + # Start with one workflow poller to avoid burst contention at startup. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MINIMUM + value: "1" + # Small initial poller count warms up gradually instead of spiking task fetches. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_INITIAL + value: "2" + # Cap workflow pollers to limit scheduling pressure and avoid over-polling. + - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM + value: "10" + # Keep at least two activity pollers so activity queues do not starve during spikes. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM + value: "2" + # Moderate initial activity pollers for faster ramp-up with controlled pressure. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_INITIAL + value: "5" + # Limit max activity pollers to preserve CPU for workflow-task completion. + - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM + value: "20" + - name: PI_WEB_API_BASE_URL value: "https://pivision.votorantimcimentos.com/piwebapi" - name: PI_WEB_API_AUTH_TYPE From ff2999dc14131c117bfbb4a9fb3133cd0e56441b Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 26 Mar 2026 14:40:58 -0300 Subject: [PATCH 34/51] SIENTIAPDE-1712 Update on_conflict parameter in FormatAndExportPrediction class to 'error' for stricter conflict handling during prediction export. --- .../workflows/sub_workflows/format_and_export_prediction.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index a9cda2a..68552b2 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -198,7 +198,7 @@ class FormatAndExportPrediction: 'table_name': input_data['table_name'], 'data': prediction, 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, - 'on_conflict': input_data.get('on_conflict', 'replace'), + 'on_conflict': input_data.get('on_conflict', 'error'), 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=retry_policy, From 8a0bb70fc80c2d52280b33531a52b98c407e6362 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 26 Mar 2026 14:56:50 -0300 Subject: [PATCH 35/51] SIENTIAPDE-1712 Refactor prediction process to use execute_activity_method for MLFlow model transformation and prediction requests, enhancing consistency in workflow execution. --- laborious/workflows/sub_workflows/prediction_process.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 3493b59..7c4d702 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -146,7 +146,7 @@ class PredictionProcess: return # Request MLFlow model transformation - transformed_data = await workflow.execute_local_activity_method( + 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, @@ -191,7 +191,7 @@ class PredictionProcess: ): return - predicted_data = await workflow.execute_local_activity_method( + predicted_data = await workflow.execute_activity_method( Activities.request_predict, { **metadata, From aec75cd33674270dcad2ba7537cdde200e0e122c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 26 Mar 2026 14:59:54 -0300 Subject: [PATCH 36/51] SIENTIAPDE-1712 Log worker runtime configuration in prepare_worker function for improved debugging and monitoring. --- laborious/worker/prepare_worker.py | 1 + 1 file changed, 1 insertion(+) diff --git a/laborious/worker/prepare_worker.py b/laborious/worker/prepare_worker.py index 07ae36f..9766af2 100644 --- a/laborious/worker/prepare_worker.py +++ b/laborious/worker/prepare_worker.py @@ -47,6 +47,7 @@ def prepare_worker( ) logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}') + logger.info(f'Worker runtime config: {local_workflow_parameters}') return Worker( temporal_client, From d71c45e61dc6a29b61e254ea60ee1612e66e4fbe Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Fri, 27 Mar 2026 14:28:13 -0300 Subject: [PATCH 37/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Add 'on_conflict' parameter to PredictionsBatch and PredictionProcess classes for improved conflict handling during predictions. --- laborious/workflows/predictions_batch.py | 1 + laborious/workflows/sub_workflows/prediction_process.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index 36e2e97..e68c215 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -117,6 +117,7 @@ class PredictionsBatch: '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), diff --git a/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 7c4d702..da4125b 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -228,6 +228,7 @@ class PredictionProcess: '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, @@ -337,6 +338,7 @@ class PredictionProcess: '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 From f84d38a837962c9cd9d5d5c63355e86d49448166 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 30 Mar 2026 08:51:41 -0300 Subject: [PATCH 38/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Implement debug logging for DataFrames in MLFlow and MLFlowRepository classes. Added a method to log DataFrame content conditionally based on row count, enhancing debugging capabilities while preventing excessive log output. --- laborious/activities/mlflow.py | 34 ++++++++++++-- .../utils/repository/model_repository.py | 46 ++++++++++++++----- 2 files changed, 64 insertions(+), 16 deletions(-) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index ad9b6ae..d1cb50a 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -42,6 +42,7 @@ class MLFlow(MinioManager): mlflow_password (str): MLFlow authentication password model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 def __init__( self, @@ -94,6 +95,32 @@ class MLFlow(MinioManager): 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 + """ + if not hasattr(data, 'shape') or not hasattr(data, 'to_csv'): + self.debug(f'{message}\n{data}', metadata) + return + + rows = data.shape[0] + if rows <= self._MAX_DEBUG_DATAFRAME_ROWS: + self.debug(f'{message}\n{data.to_csv()}', metadata) + return + + self.debug( + ( + f'{message} skipped because dataframe has {rows} rows ' + f'(max: {self._MAX_DEBUG_DATAFRAME_ROWS}). Shape: {data.shape}' + ), + metadata, + ) + @activity.defn(name='request_transform') async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload: """ @@ -133,8 +160,7 @@ class MLFlow(MinioManager): model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) - self.debug('Raw input data:', metadata) - self.debug(data.head(5).to_string(), metadata) + 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( @@ -150,7 +176,7 @@ class MLFlow(MinioManager): data['timestamp'] = data.index - self.debug(f'Processed input data: \n {data.to_csv()}', metadata) + self._debug_dataframe('Processed input data:', data, metadata) # Request transformation from MLFlow model response_data = await self.model_monitoring_repository.transform( @@ -231,7 +257,7 @@ class MLFlow(MinioManager): model_name = input_data['model_name'] model_config = input_data.get('model_config', {}) - self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata) + self._debug_dataframe('Input data for prediction:', data, metadata) # Convert numpy.nan to None for model compatibility data.replace(np.nan, None, inplace=True) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index fc965dd..d4b2ba5 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -64,6 +64,8 @@ def force_memory_release(logger: Logger): class MLFlowRepository(SientiaMonitoring): + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, host: str, @@ -94,6 +96,32 @@ class MLFlowRepository(SientiaMonitoring): 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 + """ + if not isinstance(data, pd.DataFrame): + self.debug(f'{message} {data}', metadata) + return + + rows = data.shape[0] + if rows <= self._MAX_DEBUG_DATAFRAME_ROWS: + self.debug(f'{message}\n{data.to_csv()}', metadata) + return + + self.debug( + ( + f'{message} skipped because dataframe has {rows} rows ' + f'(max: {self._MAX_DEBUG_DATAFRAME_ROWS}). Shape: {data.shape}' + ), + metadata, + ) + """ Functions related to get model registry parameters """ @@ -1168,7 +1196,7 @@ class MLFlowRepository(SientiaMonitoring): and returned in the response structure rather than propagated. """ - self.debug(f'Data received for model transformation: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model transformation:', data, metadata) # data.to_csv( # f"tmp/data_{model_name}.csv", index=True) @@ -1186,9 +1214,8 @@ class MLFlowRepository(SientiaMonitoring): metadata=metadata, ) - self.debug( - f'Data received from model transformation: {transformed_data.head(5).to_csv()}', - metadata, + self._debug_dataframe( + 'Data received from model transformation:', transformed_data, metadata ) # transformed_data.to_csv( @@ -1254,9 +1281,7 @@ class MLFlowRepository(SientiaMonitoring): input_index = data.index start_time = datetime.now() - self.debug( - f'Data received for model prediction: {data.to_dict(orient="records")}', metadata - ) + self._debug_dataframe('Data received for model prediction:', data, metadata) # data.to_csv( # f"tmp/treated_data_{model_name}.csv", index=True) @@ -1273,10 +1298,7 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self.debug( - f'Data received from model prediction: {predict_data.to_dict(orient="records")}', - metadata, - ) + self._debug_dataframe('Data received from model prediction:', predict_data, metadata) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) @@ -1348,7 +1370,7 @@ class MLFlowRepository(SientiaMonitoring): """ self.info(f'Starting model retraining workflow for {model_name}', metadata) - self.debug(f'Data received for model retraining: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model retraining:', data, metadata) target_name = model_config.get('target', None) From a8259d716a4aee85f05b905205fc5ebde7269eb7 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 30 Mar 2026 08:54:56 -0300 Subject: [PATCH 39/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Refactor debug logging for DataFrames across multiple classes. Introduced a new method to log DataFrame content conditionally based on row count in Gates, MLFlow, ModelMetrics, and MLFlowRepository classes, improving debugging capabilities while managing log output effectively. --- laborious/activities/gates.py | 30 ++++++++++--- laborious/activities/mlflow.py | 17 +++---- laborious/activities/model_metrics.py | 44 +++++++++++-------- laborious/utils/dataframe_debug.py | 34 ++++++++++++++ .../utils/repository/model_repository.py | 17 +++---- 5 files changed, 95 insertions(+), 47 deletions(-) create mode 100644 laborious/utils/dataframe_debug.py diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 12fafdf..34376b1 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -16,6 +16,7 @@ with workflow.unsafe.imports_passed_through(): 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, @@ -86,6 +87,7 @@ class Gates(MinioManager): """ minio_repository: MinioRepository | None = None + _MAX_DEBUG_DATAFRAME_ROWS = 100 def __init__( self, @@ -118,6 +120,24 @@ class Gates(MinioManager): 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]]: """ @@ -179,7 +199,7 @@ class Gates(MinioManager): filter_output = [] - self.debug(f'Input data: {data.head(5).to_string()}', metadata) + self._debug_dataframe('Input data:', data, metadata) self.debug(f'Filters: {filters}', metadata) # Apply each configured filter @@ -355,7 +375,7 @@ class Gates(MinioManager): filter_output = [] - self.debug(f'Input data:\n {data.head(5).to_string()}', metadata) + self._debug_dataframe('Input data:', data, metadata) self.debug(f'Filters: \n {filters}', metadata) for fil, config in filters.items(): @@ -544,7 +564,7 @@ class Gates(MinioManager): data = data.reset_index(drop=True) self.debug(f'Prediction store policy: {prediction_store_policy}', metadata) - self.debug(f'Prediction data: {data.head(5).to_string()}', metadata) + self._debug_dataframe('Prediction data:', data, metadata) policy_type, policy_value = self.get_prediction_store_policy( prediction_store_policy, metadata @@ -581,7 +601,7 @@ class Gates(MinioManager): data = data.reset_index(drop=True) self.info(f'Prediction formatted: {len(data)} rows', metadata) - self.debug(f'Prediction data: {data.head(5).to_string()}', metadata) + self._debug_dataframe('Prediction data:', data, metadata) return data.to_dict() @@ -696,7 +716,7 @@ class Gates(MinioManager): report['mlflow_run_id'] = update_report['mlflow_run_id'] report['mlflow_experiment_id'] = update_report['mlflow_experiment_id'] - self.debug(f'Retrain report: {report.to_csv()}', metadata) + self._debug_dataframe('Retrain report:', report, metadata) return report.to_dict() diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index d1cb50a..10d3d8d 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.utils.formatters import create_sample_dict from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload + from laborious.utils.dataframe_debug import build_dataframe_debug_message from laborious.utils.repository.minio_manager import MinioManager from laborious.utils.repository.model_repository import MLFlowRepository @@ -104,19 +105,11 @@ class MLFlow(MinioManager): - data (Any): Dataframe-like object expected to expose shape and to_csv - metadata (dict[str, Any]): Workflow metadata for contextual logging """ - if not hasattr(data, 'shape') or not hasattr(data, 'to_csv'): - self.debug(f'{message}\n{data}', metadata) - return - - rows = data.shape[0] - if rows <= self._MAX_DEBUG_DATAFRAME_ROWS: - self.debug(f'{message}\n{data.to_csv()}', metadata) - return - self.debug( - ( - f'{message} skipped because dataframe has {rows} rows ' - f'(max: {self._MAX_DEBUG_DATAFRAME_ROWS}). Shape: {data.shape}' + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, ), metadata, ) diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index 29b2848..b336fbd 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -17,6 +17,7 @@ with workflow.unsafe.imports_passed_through(): 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( @@ -26,6 +27,8 @@ warnings.filterwarnings( class ModelMetrics(SientiaMonitoring): """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 + Metrics activities for the Laborious system. This class provides activities for writing metrics to the Prometheus monitoring system. @@ -48,6 +51,24 @@ class ModelMetrics(SientiaMonitoring): 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, @@ -78,14 +99,9 @@ class ModelMetrics(SientiaMonitoring): model_analysis = ModelAnalysis(config=config) - self.debug( - f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}', - metadata, - ) + self._debug_dataframe(f'Reference data: Size {reference_data.shape}', reference_data, metadata) - self.debug( - f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', 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() @@ -142,9 +158,7 @@ class ModelMetrics(SientiaMonitoring): 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( - f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata - ) + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) return drift_df @@ -274,11 +288,7 @@ class ModelMetrics(SientiaMonitoring): drift_df['timestamp'] = drift_df['timestamp'].dt.tz_localize('UTC') drift_df['timestamp'] = drift_df['timestamp'].dt.strftime(DATETIME_FORMAT_WITH_TZ) - self.debug( - f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata - ) - - self.debug(f'Drift dataframe: {drift_df.head(5).to_string()}', metadata) + self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata) return drift_df.to_dict(orient='records') @@ -347,8 +357,6 @@ class ModelMetrics(SientiaMonitoring): data['data_size'] = data_size data['interval_minutes'] = interval_minutes - self.debug( - f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata - ) + self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata) return data.to_dict(orient='records') 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/repository/model_repository.py b/laborious/utils/repository/model_repository.py index d4b2ba5..fb404b7 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -36,6 +36,7 @@ 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' @@ -105,19 +106,11 @@ class MLFlowRepository(SientiaMonitoring): - data (Any): Dataframe-like object expected to expose shape and to_csv - metadata (dict[str, Any]): Metadata for contextual logging """ - if not isinstance(data, pd.DataFrame): - self.debug(f'{message} {data}', metadata) - return - - rows = data.shape[0] - if rows <= self._MAX_DEBUG_DATAFRAME_ROWS: - self.debug(f'{message}\n{data.to_csv()}', metadata) - return - self.debug( - ( - f'{message} skipped because dataframe has {rows} rows ' - f'(max: {self._MAX_DEBUG_DATAFRAME_ROWS}). Shape: {data.shape}' + build_dataframe_debug_message( + message=message, + data=data, + max_rows=self._MAX_DEBUG_DATAFRAME_ROWS, ), metadata, ) From 00ac77a0912f63f695cc990921e6b371535c333d Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 30 Mar 2026 09:12:27 -0300 Subject: [PATCH 40/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Implement debug logging in MinioDataFramePayload class for enhanced traceability. Added a static method for conditional logging and integrated debug statements throughout methods to capture DataFrame size estimates, upload actions, and retrieval processes, improving overall observability. --- .../utils/models/minio_dataframe_payload.py | 94 ++++++++++++++++++- 1 file changed, 89 insertions(+), 5 deletions(-) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index bb3f2fd..30f5a00 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -20,6 +20,7 @@ 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 @@ -81,6 +82,24 @@ class MinioDataFramePayload: 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.debug(message, metadata) + @classmethod def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload': """ @@ -114,7 +133,11 @@ class MinioDataFramePayload: ) @staticmethod - def estimate_size_bytes(df: DataFrame) -> int: + 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. @@ -125,9 +148,13 @@ class MinioDataFramePayload: int: Estimated size in bytes (pickle of dict representation). """ try: - return len(pickle.dumps(df.to_dict())) + size = len(pickle.dumps(df.to_dict())) except Exception: - return len(pickle.dumps(df)) + size = len(pickle.dumps(df)) + + if logger is not None: + logger.debug(f'DataFrame size: {size} bytes', metadata) + return size @staticmethod def parse_object_timestamp(object_key: str) -> datetime | None: @@ -174,6 +201,7 @@ class MinioDataFramePayload: 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. @@ -196,6 +224,11 @@ class MinioDataFramePayload: """ 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 ) @@ -203,11 +236,34 @@ class MinioDataFramePayload: if last_timestamp is None: last_timestamp = max(dataframe['timestamp'].values.tolist()) - if cls.estimate_size_bytes(dataframe) <= OFFLOAD_THRESHOLD_BYTES: + dataframe_size = cls.estimate_size_bytes(dataframe) + 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. @@ -224,6 +280,11 @@ class MinioDataFramePayload: 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, @@ -236,7 +297,10 @@ class MinioDataFramePayload: ) async def retrieve( - self, minio_repo: MinioRepository, workflow_metadata: dict[str, Any] | None = None + 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. @@ -249,13 +313,33 @@ class MinioDataFramePayload: 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 From fa885721266ffa8b9a1870e8c4635f7de5c4dc02 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 30 Mar 2026 09:24:52 -0300 Subject: [PATCH 41/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Enhance logging in Gates, MLFlow, and Storage classes by integrating logger parameter for improved traceability. This update allows for better monitoring of operations and data handling across these components. --- laborious/activities/gates.py | 1 + laborious/activities/mlflow.py | 4 ++++ laborious/activities/storage.py | 1 + laborious/utils/models/minio_dataframe_payload.py | 2 +- 4 files changed, 7 insertions(+), 1 deletion(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 34376b1..b53562e 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -522,6 +522,7 @@ class Gates(MinioManager): operation='transform', workflow_metadata=metadata, last_timestamp=payload.last_timestamp, + logger=self.logger ) @activity.defn(name='format_prediction') diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 10d3d8d..207e20e 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -197,6 +197,7 @@ class MLFlow(MinioManager): status=response_data, workflow_metadata=metadata, last_timestamp=payload.last_timestamp, + logger=self.logger, ) return await MinioDataFramePayload.from_dataframe( @@ -209,6 +210,7 @@ class MLFlow(MinioManager): 'success': True, }, last_timestamp=payload.last_timestamp, + logger=self.logger, ) @activity.defn(name='request_predict') @@ -281,6 +283,7 @@ class MLFlow(MinioManager): status=response_data, workflow_metadata=metadata, last_timestamp=payload.last_timestamp, + logger=self.logger, ) return await MinioDataFramePayload.from_dataframe( @@ -293,6 +296,7 @@ class MLFlow(MinioManager): 'success': True, }, last_timestamp=payload.last_timestamp, + logger=self.logger, ) @activity.defn(name='retrain_model') diff --git a/laborious/activities/storage.py b/laborious/activities/storage.py index 6a6207a..d64cc52 100644 --- a/laborious/activities/storage.py +++ b/laborious/activities/storage.py @@ -105,6 +105,7 @@ class Storage(Postgres, MinioManager): workflow_metadata=metadata, model_name=model_name, operation='initial', + logger=self.logger, ) @activity.defn(name='export_payload_to_postgres') diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index 30f5a00..b480ee1 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -236,7 +236,7 @@ class MinioDataFramePayload: if last_timestamp is None: last_timestamp = max(dataframe['timestamp'].values.tolist()) - dataframe_size = cls.estimate_size_bytes(dataframe) + dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger) cls._debug( logger, ( From 7c32424cf1930814c52abaf909982293cc716a2a Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Mon, 30 Mar 2026 09:32:51 -0300 Subject: [PATCH 42/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Refactor logging in MinioDataFramePayload class to utilize custom_debug method for improved clarity and consistency. Enhanced DataFrame size logging by integrating a dedicated debug method, streamlining the logging process. --- laborious/utils/models/minio_dataframe_payload.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/laborious/utils/models/minio_dataframe_payload.py b/laborious/utils/models/minio_dataframe_payload.py index b480ee1..0521f5e 100644 --- a/laborious/utils/models/minio_dataframe_payload.py +++ b/laborious/utils/models/minio_dataframe_payload.py @@ -98,7 +98,7 @@ class MinioDataFramePayload: """ if logger is None: return - logger.debug(message, metadata) + logger.custom_debug(message, metadata) @classmethod def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload': @@ -152,8 +152,11 @@ class MinioDataFramePayload: except Exception: size = len(pickle.dumps(df)) - if logger is not None: - logger.debug(f'DataFrame size: {size} bytes', metadata) + MinioDataFramePayload._debug( + logger, + f'DataFrame size: {size} bytes', + metadata, + ) return size @staticmethod From d9ec7fc4967cb903e4a67d5a590cb636f86b138b Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 31 Mar 2026 09:10:56 -0300 Subject: [PATCH 43/51] SIENTIAPDE-1712 Update values.yaml to rename SIE_MINIO_OFFLOAD_THRESHOLD_MEGABYTES to SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES and adjust its value to 0.5 for improved configuration clarity. --- values.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/values.yaml b/values.yaml index 51e96fb..88f0393 100644 --- a/values.yaml +++ b/values.yaml @@ -231,8 +231,8 @@ env: value: "sientia" - name: MINIO_RETENTION_HOURS value: "24" - - name: SIE_MINIO_OFFLOAD_THRESHOLD_MEGABYTES - value: "1.5" + - name: SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES + value: "0.5" # Temporal worker tuning for PredictionsBatch. # IMPORTANT: prefix must be PREDICTIONSBATCH_ (from class name PredictionsBatch). From a3775279e12e5f3c774e8338280eb0a8fa21bdcf Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 31 Mar 2026 15:05:40 -0300 Subject: [PATCH 44/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Refactor logging in MLFlowRepository class by removing the debug_dataframe method and replacing it with direct debug statements for improved clarity. This change enhances the logging of DataFrame content during model transformation, prediction, and retraining processes, ensuring better observability without excessive log output. --- .../utils/repository/model_repository.py | 70 ++++++------------- 1 file changed, 21 insertions(+), 49 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index fb404b7..337eda2 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -36,7 +36,6 @@ 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' @@ -65,8 +64,6 @@ def force_memory_release(logger: Logger): class MLFlowRepository(SientiaMonitoring): - _MAX_DEBUG_DATAFRAME_ROWS = 100 - def __init__( self, host: str, @@ -97,24 +94,6 @@ class MLFlowRepository(SientiaMonitoring): 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 """ @@ -459,26 +438,17 @@ class MLFlowRepository(SientiaMonitoring): raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") artifact_path = None + 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) 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' + model = model._model_impl.python_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}' - ) - - raw_model = mlflow.pyfunc.load_model(artifact_path) - model = raw_model._model_impl.python_model - 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) + self.debug(f"Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}", metadata) return model, artifact_path @@ -1155,7 +1125,7 @@ class MLFlowRepository(SientiaMonitoring): 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. @@ -1189,7 +1159,7 @@ class MLFlowRepository(SientiaMonitoring): and returned in the response structure rather than propagated. """ - self._debug_dataframe('Data received for model transformation:', data, metadata) + self.debug(f'Data received for model transformation: {data.to_csv()}', metadata) # data.to_csv( # f"tmp/data_{model_name}.csv", index=True) @@ -1207,8 +1177,9 @@ class MLFlowRepository(SientiaMonitoring): metadata=metadata, ) - self._debug_dataframe( - 'Data received from model transformation:', transformed_data, metadata + self.debug( + f'Data received from model transformation: {transformed_data.head(5).to_csv()}', + metadata, ) # transformed_data.to_csv( @@ -1216,7 +1187,7 @@ class MLFlowRepository(SientiaMonitoring): transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return {'success': True, 'content': transformed_data} + return {'success': True, 'content': transformed_data.to_dict()} except Exception as e: return { @@ -1274,7 +1245,9 @@ class MLFlowRepository(SientiaMonitoring): input_index = data.index start_time = datetime.now() - self._debug_dataframe('Data received for model prediction:', data, metadata) + self.debug( + f'Data received for model prediction: {data.to_dict(orient="records")}', metadata + ) # data.to_csv( # f"tmp/treated_data_{model_name}.csv", index=True) @@ -1291,17 +1264,16 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self._debug_dataframe('Data received from model prediction:', predict_data, metadata) + self.debug( + f'Data received from model prediction: {predict_data.to_dict(orient="records")}', + 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) @@ -1309,7 +1281,7 @@ class MLFlowRepository(SientiaMonitoring): predict_data.index = input_index predict_data['response_time'] = (end_time - start_time).total_seconds() - return {'success': True, 'content': predict_data} + return {'success': True, 'content': predict_data.to_dict()} except Exception as e: return { @@ -1363,7 +1335,7 @@ class MLFlowRepository(SientiaMonitoring): """ self.info(f'Starting model retraining workflow for {model_name}', metadata) - self._debug_dataframe('Data received for model retraining:', data, metadata) + self.debug(f'Data received for model retraining: {data.to_csv()}', metadata) target_name = model_config.get('target', None) From e6fee73cce28b8af65339cf123e12024dd0aa1c8 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 31 Mar 2026 15:13:05 -0300 Subject: [PATCH 45/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Enhance logging in MLFlowRepository class by introducing a dedicated _debug_dataframe method for conditional logging of DataFrame content. This update improves observability during model transformation, prediction, and retraining processes while managing log output effectively. --- .../utils/repository/model_repository.py | 70 +++++++++++++------ 1 file changed, 50 insertions(+), 20 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 337eda2..581e57f 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -36,6 +36,7 @@ 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' @@ -64,6 +65,8 @@ def force_memory_release(logger: Logger): class MLFlowRepository(SientiaMonitoring): + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, host: str, @@ -94,6 +97,24 @@ class MLFlowRepository(SientiaMonitoring): 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 """ @@ -438,17 +459,28 @@ class MLFlowRepository(SientiaMonitoring): raise ValueError("Invalid model_type. Use 'predict' or 'transform'.") artifact_path = None - 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) if load_wrapper: self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}') - model = model._model_impl.python_model + 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}' + ) + + raw_model = mlflow.pyfunc.load_model(artifact_path) + model = raw_model._model_impl.python_model 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 @@ -1125,7 +1157,7 @@ class MLFlowRepository(SientiaMonitoring): 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. @@ -1159,7 +1191,7 @@ class MLFlowRepository(SientiaMonitoring): and returned in the response structure rather than propagated. """ - self.debug(f'Data received for model transformation: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model transformation:', data, metadata) # data.to_csv( # f"tmp/data_{model_name}.csv", index=True) @@ -1177,9 +1209,8 @@ class MLFlowRepository(SientiaMonitoring): metadata=metadata, ) - self.debug( - f'Data received from model transformation: {transformed_data.head(5).to_csv()}', - metadata, + self._debug_dataframe( + 'Data received from model transformation:', transformed_data, metadata ) # transformed_data.to_csv( @@ -1187,7 +1218,7 @@ class MLFlowRepository(SientiaMonitoring): transformed_data = self.detect_and_parse_datetime_index(transformed_data, metadata) - return {'success': True, 'content': transformed_data.to_dict()} + return {'success': True, 'content': transformed_data} except Exception as e: return { @@ -1245,9 +1276,7 @@ class MLFlowRepository(SientiaMonitoring): input_index = data.index start_time = datetime.now() - self.debug( - f'Data received for model prediction: {data.to_dict(orient="records")}', metadata - ) + self._debug_dataframe('Data received for model prediction:', data, metadata) # data.to_csv( # f"tmp/treated_data_{model_name}.csv", index=True) @@ -1264,16 +1293,17 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self.debug( - f'Data received from model prediction: {predict_data.to_dict(orient="records")}', - metadata, - ) + 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) @@ -1281,7 +1311,7 @@ class MLFlowRepository(SientiaMonitoring): predict_data.index = input_index predict_data['response_time'] = (end_time - start_time).total_seconds() - return {'success': True, 'content': predict_data.to_dict()} + return {'success': True, 'content': predict_data} except Exception as e: return { @@ -1335,7 +1365,7 @@ class MLFlowRepository(SientiaMonitoring): """ self.info(f'Starting model retraining workflow for {model_name}', metadata) - self.debug(f'Data received for model retraining: {data.to_csv()}', metadata) + self._debug_dataframe('Data received for model retraining:', data, metadata) target_name = model_config.get('target', None) From 59dd6078b09445cc88b384dbefa48065b2975b21 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 31 Mar 2026 16:21:32 -0300 Subject: [PATCH 46/51] SIENTIAPDE-1712 Update values.yaml to adjust workflow and activity poller configurations for improved performance, increasing minimum, initial, and maximum values. Additionally, add new minimal retrain parameters to enhance retraining capabilities. --- values.yaml | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/values.yaml b/values.yaml index 88f0393..31dd4fa 100644 --- a/values.yaml +++ b/values.yaml @@ -250,22 +250,41 @@ env: value: "200" # Start with one workflow poller to avoid burst contention at startup. - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MINIMUM - value: "1" + value: "3" # Small initial poller count warms up gradually instead of spiking task fetches. - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_INITIAL - value: "2" + value: "5" # Cap workflow pollers to limit scheduling pressure and avoid over-polling. - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM - value: "10" + value: "15" # Keep at least two activity pollers so activity queues do not starve during spikes. - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM - value: "2" + value: "3" # Moderate initial activity pollers for faster ramp-up with controlled pressure. - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_INITIAL - value: "5" + value: "10" # Limit max activity pollers to preserve CPU for workflow-task completion. - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM - value: "20" + value: "30" + + - name: MINIMALRETRAIN_MAX_CONCURRENT_ACTIVITIES + value: "1" + - name: MINIMALRETRAIN_MAX_CONCURRENT_LOCAL_ACTIVITIES + value: "1" + - name: MINIMALRETRAIN_MAX_CACHED_WORKFLOWS + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_MINIMUM + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_INITIAL + value: "1" + - name: MINIMALRETRAIN_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_INITIAL + value: "1" + - name: MINIMALRETRAIN_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM + value: "1" - name: PI_WEB_API_BASE_URL value: "https://pivision.votorantimcimentos.com/piwebapi" From b9fe4604f7f144075514def4bfb7d8da1475cc29 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 1 Apr 2026 09:30:55 -0300 Subject: [PATCH 47/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Convert confidence_value to float in API class to ensure consistent data type for value assignment in activity monitoring. --- laborious/activities/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/laborious/activities/api.py b/laborious/activities/api.py index d91b0e3..0e4114f 100644 --- a/laborious/activities/api.py +++ b/laborious/activities/api.py @@ -279,7 +279,7 @@ class API(SientiaMonitoring): web_ids=confidence_tags, value={ 'Timestamp': data.head(1)['timestamp'].values[0], - 'Value': confidence_value, + 'Value': float(confidence_value), }, metadata=metadata, ) From 381856f5caca828dc1d520d556613644c9cf7fb0 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 1 Apr 2026 13:53:23 -0300 Subject: [PATCH 48/51] SIENTIAPDE-1712 SIENTIAPDE-1712 Enhance logging across various classes by adding logger parameters and improving debug statements. This update includes adjustments in Gates, MLFlow, ModelMetrics, and MLFlowRepository classes for better traceability and observability during operations. --- laborious/activities/gates.py | 2 +- laborious/activities/mlflow.py | 3 +- laborious/activities/model_metrics.py | 8 +-- .../utils/repository/model_repository.py | 8 ++- tests/laborious/activities/test_gates.py | 4 ++ tests/laborious/activities/test_mlflow.py | 2 + .../test_format_and_export_prediction.py | 24 ++++++--- .../subworkflows/test_prediction_process.py | 53 ++++++++++++------- .../workflows/test_predictions_batch.py | 1 + 9 files changed, 74 insertions(+), 31 deletions(-) diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index b53562e..ad88ca1 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -522,7 +522,7 @@ class Gates(MinioManager): operation='transform', workflow_metadata=metadata, last_timestamp=payload.last_timestamp, - logger=self.logger + logger=self.logger, ) @activity.defn(name='format_prediction') diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 207e20e..f4d42a3 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -19,8 +19,8 @@ with workflow.unsafe.imports_passed_through(): ) from sientia_do.utils.formatters import create_sample_dict - from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload 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 @@ -43,6 +43,7 @@ class MLFlow(MinioManager): mlflow_password (str): MLFlow authentication password model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 def __init__( diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index b336fbd..79b140c 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -27,13 +27,13 @@ warnings.filterwarnings( class ModelMetrics(SientiaMonitoring): """ - _MAX_DEBUG_DATAFRAME_ROWS = 100 - 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, @@ -99,7 +99,9 @@ class ModelMetrics(SientiaMonitoring): model_analysis = ModelAnalysis(config=config) - self._debug_dataframe(f'Reference data: Size {reference_data.shape}', reference_data, metadata) + 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) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 581e57f..b4fdad5 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -474,7 +474,9 @@ class MLFlowRepository(SientiaMonitoring): raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model - self.debug(f"Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}", metadata) + 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) @@ -1293,7 +1295,9 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self._debug_dataframe('Data received from model prediction:', predict_data, metadata) + self._debug_dataframe( + 'Data received from model prediction:', predict_data, metadata + ) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 3d072b9..c144068 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -541,6 +541,7 @@ async def test_format_prediction_no_timestamp(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:1', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -580,6 +581,7 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'erl:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -619,6 +621,7 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -655,6 +658,7 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', + 'timestamp': '2023-05-26 11:12:27', } gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1)) diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index daaebab..5d242e5 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -181,6 +181,7 @@ async def test_request_transform_failure(mock_from_dataframe, mlflow): status=transform_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, + logger=mlflow.logger, ) assert response_data == mock_from_dataframe.return_value @@ -252,6 +253,7 @@ async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, ml status=predict_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, + logger=mlflow.logger, ) assert response_data == mock_from_dataframe.return_value diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index 6e8817f..4fbc62b 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -93,6 +93,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -100,7 +101,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -243,6 +245,7 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -250,7 +253,8 @@ async def test_run_none_path_flag_with_transformed_data( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -349,14 +353,16 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, - **metadata, 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -456,6 +462,7 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -463,7 +470,8 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -560,6 +568,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -567,7 +576,8 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -648,6 +658,7 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -655,7 +666,8 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 60f3744..a018504 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -51,14 +51,17 @@ async def test_run(workflow_mock, prediction_process): } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict # mlflow_response_gate (predict) ('continue', 0.95, 'Error'), ] @@ -67,7 +70,7 @@ async def test_run(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 6 + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -83,7 +86,7 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -130,7 +133,7 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, @@ -166,6 +169,7 @@ async def test_run(workflow_mock, prediction_process): 'subworkflow.format_and_export_prediction', { 'metadata': metadata, + 'on_conflict': 'error', 'path_flag': 'continue', 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, @@ -266,9 +270,12 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('repeat', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data ('continue', 0.95, 'Error'), # mlflow_response_gate (transform) ] @@ -276,7 +283,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 3 + assert workflow_mock.execute_local_activity_method.call_count == 2 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -292,7 +299,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -353,9 +360,12 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) @@ -366,7 +376,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 4 + assert workflow_mock.execute_local_activity_method.call_count == 3 workflow_mock.execute_local_activity_method.assert_has_calls( [ @@ -383,7 +393,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -460,14 +470,17 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict ('continue', 0.95, 'Error'), # mlflow_response_gate (predict) ] @@ -475,7 +488,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 6 + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -491,7 +504,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -538,7 +551,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, @@ -727,6 +740,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'confidence_tags': {}, }, 'prediction_store_policy': prediction_store_policy, + 'on_conflict': 'error', }, ) @@ -806,12 +820,15 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): 'prediction_store_policy': 'lts:1', } + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'ok'), - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, ('continue', 0.95, ''), ('continue', 0.95, ''), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, ('continue', 0.95, ''), ] diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index de52ad8..14280d6 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -95,6 +95,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch '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), From 5398ac6bc2e11eacd7d910336c54d1b5a1a4f7e3 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 1 Apr 2026 16:53:12 -0300 Subject: [PATCH 49/51] SIENTIAPDE-1712 Update sientia-mlops-library dependency to version 0.41.0 in requirements.txt for improved functionality and features. --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 7cb8e5b..3f33de0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,7 +4,7 @@ sqlalchemy asyncua redis git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.10.4 -git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.40.7 +git+ssh://git@github.com/Aignosi/sientia-mlops-library.git@0.41.0 prometheus-client botocore boto3 From 5e883d45302b29a475e82d9dbcfc0b2a39bab9cf Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 2 Apr 2026 10:06:39 -0300 Subject: [PATCH 50/51] SIENTIAPDE-1712 Add scipy version 1.13.0 to requirements.txt for enhanced scientific computing capabilities. --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 3f33de0..381aebb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,4 +14,5 @@ kaleido hyperopt shap pycurl +scipy==1.13.0 scikit-learn==1.5.2 \ No newline at end of file From 25430b5b261b39649ed081d87ae90b02e7e97f8c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 2 Apr 2026 10:07:05 -0300 Subject: [PATCH 51/51] SIENTIAPDE-1712 Update scipy version constraint in requirements.txt to allow versions below 1.14.0 for compatibility with existing dependencies. --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 381aebb..6c5c242 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,5 +14,5 @@ kaleido hyperopt shap pycurl -scipy==1.13.0 +scipy<1.14.0 scikit-learn==1.5.2 \ No newline at end of file