Merge pull request #38 from Aignosi/feature/SIENTIAPDE-1712

Feat/Refactor: MinIO Offload, Data Tracking, and Drift Detection Workflows

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
+1. **Input Data Gate**: Applies configured filters for data quality validation
-2. **Input Data Gate**: Applies configured filters for data quality validation
+2. **Path Decision**: Determines processing path based on filter results
-3. **Path Decision**: Determines processing path based on filter results
+3. **MLFlow Transform**: Requests data transformation using MLFlow models
-4. **MLFlow Transform**: Requests data transformation using MLFlow models
+4. **Response Validation**: Filters transform responses for quality assurance
-5. **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)]
+    B -.-> MLFlow[MLFlow]
-    C -.-> MLFlow[MLFlow]
+    E -.-> MLFlow[MLFlow]
-    F -.-> MLFlow[MLFlow]
+    H -.-> MinIO[(MinIO)]
-    G -.-> Filters[MLFlow Filters]
 ```
 
 ### 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/laborious/activities/
-pytest tests/workflow/
+pytest tests/laborious/workflows/
 ```
 
 ### Manual Application Execution
@@ -569,18 +669,7 @@ pip install -r requirements-dev.txt
 
 ### Complete Validation
 
-Option 1 (recommended):
+Run each validation step individually:
-```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):
 ```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
+├── conftest.py                     # Global fixtures and env setup
-│   ├── test_gates.py    # Data quality gates and formatting tests
+├── laborious/
-│   ├── test_mlflow.py   # MLFlow operations and reference data tests
+│   ├── activities/                 # Activity implementation tests
-│   └── ...              # Other activity tests
+│   │   ├── test_activities.py      # Activities aggregator tests
-├── workflows/           # Workflow orchestration tests
+│   │   ├── test_gates.py           # Data quality gates and formatting tests
-│   └── subworkflows/    # Sub-workflow tests
+│   │   ├── test_mlflow.py          # MLFlow operations and reference data tests
-│       └── test_format_and_export_prediction.py  # Export workflow tests
+│   │   ├── test_storage.py         # Storage and MinIO offload tests
-├── utils/               # Utility function tests
+│   │   ├── test_model_metrics.py   # Drift and simple metrics tests
-└── integration/         # End-to-end workflow 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/laborious/activities/test_gates.py
-pytest tests/workflow/test_predictions_batch.py
+pytest tests/laborious/workflows/test_predictions_batch.py
 ```
 
 ## 📊 Monitoring and Metrics
@@ -716,10 +824,85 @@ 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 |
+| `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 |
 
 
 
 
+### 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_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.
+
+#### 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://<bucket>/<...>` string)
+  - `data` is omitted / set to `None`.
+
+When an activity needs pandas operations, it resolves references using:
+- `MinioDataFramePayload.retrieve(minio_repo)` — downloads from MinIO or returns inline data as a DataFrame
+
+#### 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:
@@ -882,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
 ```
@@ -1004,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.
+**Note**: The Laborious system is designed for production use in industrial ML environments. Ensure proper security configuration and network isolation for production deployments.

e2e/conftest.py

@@ -3,11 +3,13 @@ Pytest configuration and fixtures for E2E tests.
 """
 
 from unittest.mock import AsyncMock, MagicMock, patch
+from io import BytesIO
 
 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
@@ -25,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():
     """
@@ -186,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):
     """
@@ -216,11 +243,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
+    # sientia_do MinioRepository API
-    mock_repo.put_parquet_from_dataframe = AsyncMock(return_value='test-object-key')
+    mock_repo.bucket = 'test-bucket'
-    mock_repo.get_parquet_as_dataframe = AsyncMock(return_value=pd.DataFrame())
+    mock_repo.upload_file = AsyncMock(
-    mock_repo.minio_bucket = 'test-bucket'
+        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
 
@@ -229,15 +273,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={
+    async def _write_value(web_ids, value, metadata=None, **kwargs):
-            'Items': [
+        """
-                {
+        Mirror successful PI writes: one response item per requested web_id.
-                    'WebId': 'web_id_1'
+
-                }
+        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
 
@@ -248,7 +294,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
@@ -261,7 +307,8 @@ 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):
+    # 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
@@ -389,11 +436,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',
-            'region_name': 'us-east-1',
             'default_bucket': 'test-bucket',
+            'retention_hours': 24,
+            'secure': False,
         },
         opc_config={},
         pi_web_api_config={
@@ -416,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."""
@@ -431,22 +562,18 @@ async def temporal_worker(temporal_test_env, test_activities):
         temporal_test_env.client,
         task_queue='test-queue',
         workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
-        activities=[
+        activities=_worker_activity_list(test_activities),
-            test_activities.load_custom_query,
+    ) as worker:
-            test_activities.get_last_timestamp,
+        yield worker
-            test_activities.input_gate,
+
-            test_activities.request_transform,
+
-            test_activities.mlflow_response_gate,
+@pytest_asyncio.fixture(scope='function')
-            test_activities.mlflow_content_gate,
+async def temporal_worker_real_minio(temporal_test_env, test_activities_real_minio):
-            test_activities.request_predict,
+    """Temporal worker backed by Activities using real MinIO testcontainer."""
-            test_activities.repeat_last_prediction,
+    async with Worker(
-            test_activities.format_prediction,
+        temporal_test_env.client,
-            test_activities.format_transformed_data,
+        task_queue='test-queue',
-            test_activities.format_default_prediction,
+        workflows=[PredictionsBatch, PredictionProcess, FormatAndExportPrediction],
-            test_activities.write_pi_web_api_data,
+        activities=_worker_activity_list(test_activities_real_minio),
-            test_activities.write_opc_data,
-            test_activities.export_data_to_postgres,
-            test_activities.write_metrics,
-        ],
     ) as worker:
         yield worker

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()}'

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
+- `write_pi_web_api_data` surfaces the failure (exception handled in the activity layer)
-- Notification sent
+- Notification may be sent
-- Workflow fails after retries
+- Workflow **completes** (does not fail)
-- PostgreSQL export may not execute (depends on execution order)
+- 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
+- PI Web API error notification sent (when applicable)
-- Workflow fails
+- Workflow completes
-- May impact subsequent exports
+- 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
+- `write_opc_data` reports failure without aborting the workflow
-- Notification sent
+- Notification may be sent
-- Workflow fails after retries
+- 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
+- OPC error notification sent (when applicable)
-- Workflow fails
+- Workflow completes
-- PostgreSQL export may not execute
+- 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

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'

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

e2e/test_predictions_batch_format_export.py

@@ -2,17 +2,17 @@
 End-to-end tests for PredictionsBatch workflow - Format and Export scenarios.
 """
 
-import asyncio
+from decimal import Decimal
-from datetime import datetime
+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
 
@@ -25,13 +25,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': {},
@@ -55,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
@@ -129,35 +65,31 @@ async def test_scenario_3_1_1_default_prediction_export(
     postgres_engine,
 ):
     """
-    Scenario 3.1.1: Default Prediction Export
+    Scenario 3.1.1: Default prediction export (non-None path_flag).
-    
+
-    Description:
+    Triggers input_gate CONTINUE via SPECIFIC_VARIABLES_NULL_VALUES so
-    Error prediction path creates default prediction.
+    PredictionProcess calls FormatAndExportPrediction with path_flag set.
-    
+    That workflow uses format_default_prediction (not format_prediction) and
-    Expected Behavior:
+    skips format_transformed_data / transform Postgres export.
-    - format_default_prediction called instead of format_prediction
+
-    - Default prediction created with error metadata
+    Optional PI Web API and OPC outputs still run when configured.
-    - 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
     """
     client = temporal_test_env.client
 
     model_id = 311
 
-    print("\n[TEST] 1. Inserting test data...")
+    with postgres_engine.begin() as conn:
-    insert_sample_data(postgres_engine, model_id, [23.5, 78.2])
+        conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}"))
-    print("[TEST] ✓ Data inserted successfully")
+        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'},
@@ -178,10 +110,9 @@ async def test_scenario_3_1_1_default_prediction_export(
         }
     }
 
-    print("\n[TEST] 2. Starting workflow that should create default prediction...")
+    wid = make_workflow_id('test-default-prediction')
-    workflow_id = f'test-default-prediction-{datetime.now().timestamp()}'
+
-    
+    await start_and_await_workflow(client, PredictionsBatch.run, input_data, wid)
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     test_activities.pi_web_api_client.write_value.assert_has_calls(
         [
@@ -189,9 +120,8 @@ 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,
                 },
-                endpoint='test_endpoint',
                 metadata={
                     'model_id': 311,
                     'model_name': 'test_model',
@@ -203,9 +133,8 @@ 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,
                 },
-                endpoint='test_endpoint',
                 metadata={
                     'model_id': 311,
                     'model_name': 'test_model',
@@ -217,28 +146,51 @@ 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,
+            call(
-            {
+                'addr_1',
-                'model_id': 311,
+                0,
-                'model_name': 'test_model',
+                'float',
-                'schedule_name': 'test-schedule',
+                ANY,
-                'workflow_name': 'predictions_batch',
+                {
-            }),
+                    'model_id': 311,
-            call('addr_2', 0, 'float', ANY,
+                    'model_name': 'test_model',
-            {
+                    'schedule_name': 'test-schedule',
-                'model_id': 311,
+                    'workflow_name': 'predictions_batch',
-                '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,12 +245,12 @@ 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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-opc-only-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-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,
             {
@@ -323,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
@@ -357,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'] = {
@@ -369,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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-pi-api-only-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-only')
-    
+    )
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     test_activities.pi_web_api_client.write_value.assert_has_calls(
         [
@@ -382,7 +328,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 +341,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,11 +352,11 @@ 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)
 
-    print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -445,25 +389,22 @@ 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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-no-optional-outputs-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-no-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)
 
-    print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -494,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
@@ -522,10 +461,9 @@ async def test_scenario_3_1_5_export_without_transformed_data(
         }
     }
 
-    print("\n[TEST] 2. Starting workflow without transformed data export...")
+    await start_and_await_workflow(
-    workflow_id = f'test-no-transform-export-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-no-transform-export')
-    
+    )
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     test_activities.pi_web_api_client.write_value.assert_has_calls(
         [
@@ -535,7 +473,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 +486,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 +497,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,
             {
@@ -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:
+    Export failure is handled inside the activity; there is no retry loop. The
-    PI Web API export fails.
+    workflow completes and PostgreSQL stores prediction_confidence 13 and the
-    
+    error message in comments.
-    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
     """
     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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-pi-api-error-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-error')
-    
+    )
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     assert_prediction(
         postgres_engine, model_id,
         prediction_confidence=13,
         comments='PI Web API service unavailable',
     )
-    
+    assert notification_inserts.call_count >= 1
-    print("\n[TEST] ✓ All assertions passed!")
+
 
 
 @pytest.mark.asyncio
@@ -673,29 +597,18 @@ async def test_scenario_3_2_2_opc_write_error(
 ):
     """
     Scenario 3.2.2: OPC Write Error
-    
+
-    Description:
+    OPC failure is reported without failing the workflow; there is no retry
-    OPC server write fails.
+    loop. PostgreSQL stores prediction_confidence 12 and OPC error comments.
-    
-    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
     """
     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")
 
-    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',
@@ -724,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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-opc-error-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-opc-error')
-    
+    )
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     assert_prediction(
         postgres_engine, model_id,
@@ -735,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
@@ -748,51 +659,24 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error(
 ):
     """
     Scenario 3.2.3: PI Web API Partial Write Error
-    
+
-    Description:
+    Partial PI write: confidence 13, descriptive comments, workflow completes
-    Two prediction tags attempt to be written to PI Web API, but only one succeeds.
+    without an activity retry loop.
-    
-    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
     """
     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=[
+    test_activities.pi_web_api_client.write_value = AsyncMock(
-        {
+        side_effect=[
-            'Items': [
+            # Prediction batch: two web_ids requested, only one acknowledged.
-                {
+            [{'WebId': 'web_id_1', 'Errors': []}],
-                    'WebId': 'web_id_1',
+            # Confidence write succeeds.
-                    'Errors': [],
+            [{'WebId': 'web_id_2', 'Errors': []}],
-                },
+        ]
-            ]
+    )
-        },
-        Exception('Tag write failed'),
-        {
-            'Items': [
-                {
-                    'WebId': 'web_id_2',
-                    'Errors': [],
-                },
-            ]
-        },
-    ])
 
     input_data = get_base_input_data(model_id)
     input_data['pi_web_api_output_config'] = {
@@ -815,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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-pi-api-partial-error-{datetime.now().timestamp()}'
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-pi-api-partial-error')
-    
+    )
-    await start_and_await_workflow(client, input_data, workflow_id)
 
     assert_prediction(
         postgres_engine, model_id,
@@ -826,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!")

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 with SQL load, MLflow mocks, Postgres predictions and transforms."""
-    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)
-    """
     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'))
-        conn.execute(text("DELETE FROM predictions_schema.laborious_data WHERE model_id = 123"))
-        
-        # Insert sensor data that the workflow will query
         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': {
@@ -100,13 +54,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': {},
@@ -121,77 +75,44 @@ async def test_scenario_1_1_1_happy_path_complete_success(
         'datetime_columns': ['timestamp', 'created_at'],
     }
 
-    # Start workflow
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow...")
+        client,
-    workflow_id = f'test-predictions-batch-{datetime.now().timestamp()}'
-    print(f"[TEST] Workflow ID: {workflow_id}")
-    
-    handle = await client.start_workflow(
         PredictionsBatch.run,
         input_data,
-        id=workflow_id,
+        make_workflow_id('test-predictions-batch'),
-        task_queue='test-queue',
     )
-    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()
-
+        assert len(prediction_rows) == 1
-        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
         row = prediction_rows[0]
-        print(f"[TEST] Prediction: {row}")
+        assert row[0] == 123
-        assert row[0] == 123, f"Expected model_id=123, got {row[0]}"
+        assert row[1] == 0.5
-        assert row[1] == 0.5, f"Expected prediction=0.5, got {row[1]}"
+        assert row[2] == 0, f'Expected prediction_confidence=0, got {row[2]}'
-        assert row[2] == 0, f"Expected prediction_confidence=0.9, got {row[2]}"
+        assert row[3] is not None
-        assert row[3] is not None, f"Expected response_time=0.1, got {row[3]}"
+        assert row[4] == 'Good'
-        assert row[4] == 'Good', f"Expected prediction_status='Good', got {row[4]}"
+        assert row[5] == ''
-        assert row[5] == '', f"Expected comments='', got {row[5]}"
-        print("[TEST] ✓ Prediction data verified")
 
-        # 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
-        assert len(transformed_rows) == 2, "Expected two transformed data records"
+        assert transformed_rows[0][0] == 123
-        row_1 = transformed_rows[0]
+        assert transformed_rows[0][1] == 'feature_1'
-        print(f"[TEST] Transformed data: {row_1}")
+        assert float(transformed_rows[0][2]) == 0.234
-        assert row_1[0] == 123, f"Expected model_id=123, got {row_1[0]}"
+        assert transformed_rows[1][0] == 123
-        assert row_1[1] == 'feature_1', f"Expected variable='sensor_1', got {row_1[1]}"
+        assert transformed_rows[1][1] == 'feature_2'
-        assert float(row_1[2]) == 0.234, f"Expected value=0.234, got {row_1[2]}"
+        assert float(transformed_rows[1][2]) == 0.783
-        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!")
 
 
 @pytest.mark.asyncio
@@ -202,23 +123,7 @@ async def test_scenario_1_2_1_sql_query_execution_error(
     test_activities: Activities,
     postgres_engine,
 ):
-    """
+    """Invalid SQL: workflow may complete with early exit; no prediction rows."""
-    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
-    """
     client = temporal_test_env.client
 
     input_data = {
@@ -233,18 +138,18 @@ 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',
         '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': {},
@@ -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...")
+    await start_and_await_workflow(
-    workflow_id = f'test-sql-error-{datetime.now().timestamp()}'
+        client,
-    
-    handle = await client.start_workflow(
         PredictionsBatch.run,
         input_data,
-        id=workflow_id,
+        make_workflow_id('test-sql-error'),
-        task_queue='test-queue',
     )
-    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(
+        count = conn.execute(
-            text("SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128")
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 128')
-        )
+        ).scalar()
-        count = result_query.scalar()
+    assert count == 0
-        assert count == 0, f"Expected no predictions, but found {count} records"
-    
-    print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -296,24 +185,9 @@ async def test_scenario_1_2_2_missing_required_parameters(
     test_activities: Activities,
     postgres_engine,
 ):
-    """
+    """Missing query: workflow does not produce predictions and is terminated explicitly."""
-    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
-    """
     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...")
+    # Let Temporal process a few workflow tasks; for this case, result() can hang.
-    try:
+    await asyncio.sleep(2.0)
-        await asyncio.wait_for(handle.result(), timeout=60.0)
+
-        pytest.fail("Expected workflow to fail, but it completed successfully")
+    with postgres_engine.connect() as conn:
-    except Exception as e:
+        count = conn.execute(
-        print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 129')
-    
+        ).scalar()
-    print("\n[TEST] ✓ All assertions passed!")
+    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,
 ):
-    """
+    """Invalid datetime column: no predictions persisted; workflow terminated after validation."""
-    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
-    """
     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"))
+        conn.execute(text('DELETE FROM predictions_schema.laborious_data WHERE model_id = 130'))
-        
+        conn.execute(
-        insert_sql = """
+            text(
+                """
             INSERT INTO predictions_schema.laborious_data (model_id, variable, value, timestamp, created_at)
-            VALUES 
+            VALUES (130, 'sensor_1', 23.5, '2024-01-01 12:00:00+00:00', '2024-01-01 12:00:00+00:00')
-                (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': {
@@ -407,13 +263,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': {},
@@ -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...")
+    # Let Temporal process and surface the failure path internally.
-    try:
+    await asyncio.sleep(2.0)
-        await asyncio.wait_for(handle.result(), timeout=60.0)
+
-        # Workflow may complete or fail depending on error handling
+    with postgres_engine.connect() as conn:
-        print("[TEST] ✓ Workflow completed (may have handled error gracefully)")
+        count = conn.execute(
-    except Exception as e:
+            text('SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = 130')
-        print(f"[TEST] ✓ Workflow failed as expected: {type(e).__name__}")
+        ).scalar()
-    
+    assert count == 0
-    print("\n[TEST] ✓ Test completed!")
+
+    await handle.terminate('expected failure path in e2e test (invalid datetime column)')

e2e/test_predictions_batch_prediction_process.py

@@ -2,55 +2,62 @@
 End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios.
 """
 
-import asyncio
-from datetime import datetime
 from decimal import Decimal
 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',
+    'schedule_name': 'test-schedule',
-        'model_name': 'test_model',
+    'model_name': 'test_model',
-        'model_id': 201,
+    'model_id': 201,
-        'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
+    'query': 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = 201',
-        'schema': 'predictions_schema',
+    'schema': 'predictions_schema',
-        'table_name': 'predictions',
+    'table_name': 'predictions',
-        'transform_table_name': 'transformed_data',
+    'transform_table_name': 'transformed_data',
-        'input_filters': {
+    'input_filters': {
-            'SPECIFIC_VARIABLES_NULL_VALUES': {
+        'SPECIFIC_VARIABLES_NULL_VALUES': {
-                'policy': 'CONTINUE',  # Continue despite issues, not STOP
+            'POLICY': 'CONTINUE',
-                'config': {'variables': ['sensor_1']},
+            'CONFIG': {'variables': ['sensor_1']},
-            },
         },
-        'mlflow_transform_filters': {
+    },
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
+    'mlflow_transform_filters': {
-        },
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
-        'mlflow_predict_filters': {
+    },
-            'API_ERROR': {'policy': 'STOP', 'config': {}},
+    'mlflow_predict_filters': {
-        },
+        'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}},
-        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
+    },
-        'opc_output_config': {},
+    'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
-        'pi_web_api_output_config': {},
+    'opc_output_config': {},
-        'save_transform': True,
+    'pi_web_api_output_config': {},
-        'prediction_store_policy': 'lts:1',
+    'save_transform': True,
-        'model_config': {
+    'prediction_store_policy': 'lts:1',
-            'retention_minutes': 0,
+    'model_config': {
-            'transform_flavor': 'sklearn',
+        'retention_minutes': 0,
-            'predict_flavor': 'sklearn',
+        'transform_flavor': 'sklearn',
-        },
+        'predict_flavor': 'sklearn',
-        'datetime_columns': ['timestamp', 'created_at'],
+    },
-    }
+    '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 {
@@ -59,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[tuple]):
-    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}"))
+        conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}'))
-        
-        # Insert data with some null values (quality issue)
-
         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')
-        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 = 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]}"
 
 
-def assert_stop(postgres_engine, model_id):
+@pytest.fixture
-    print("\n[TEST] 4. Verifying no predictions were created...")
+def bad_data_model(patch_mlflow):
-    with postgres_engine.connect() as conn:
+    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad data model')))
-        result_query = conn.execute(
+    patch_mlflow.sklearn.load_model = MagicMock(return_value=model)
-            text(f"SELECT COUNT(*) FROM predictions_schema.predictions WHERE model_id = {model_id}")
+    return model
-        )
-        count = result_query.scalar()
-        assert count == 0, f"Expected no predictions, but found {count} records"
 
-def assert_repeat(postgres_engine, model_id, last_prediction: list):
+
-    print("\n[TEST] 4. Verifying prediction was repeated...")
+@pytest.fixture
-    with postgres_engine.connect() as conn:
+def bad_predict_model(patch_mlflow, mock_mlflow_models):
-        result_query = conn.execute(
+    model = MagicMock(predict=MagicMock(side_effect=Exception('Bad predict model')))
-            text(f'SELECT model_id, prediction, prediction_confidence, prediction_status FROM predictions_schema.predictions WHERE model_id = {model_id}')
+
-        )
+    def mock_sklearn_load_model(model_uri):
-        prediction_rows = result_query.fetchall()
+        if 'data_model' in model_uri or 'transform' in model_uri.lower():
-        print(prediction_rows)
+            return mock_mlflow_models['transform_model']
-        assert len(prediction_rows) == 2, "Expected two prediction records"
+        return model
-        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}"
+    patch_mlflow.sklearn = MagicMock()
+    patch_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model)
+    return model
 
 
 @pytest.mark.asyncio
@@ -159,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,
 ):
-    """
+    """Input gate CONTINUE: export default prediction; MLflow transform/predict not used."""
-    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
-    """
     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)
-
+    await start_and_await_workflow(
-
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-continue-policy')
-    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)
     assert_continue(postgres_engine, model_id)
-    
+    mock_mlflow_models['transform_model'].predict.assert_not_called()
-    print("\n[TEST] ✓ All assertions passed!")
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -207,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,
 ):
-    """
+    """Input gate STOP: no export, no MLflow."""
-    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
-    """
     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'
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'STOP'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should stop at input gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-stop')
-    workflow_id = f'test-input-stop-{datetime.now().timestamp()}'
+    )
-    
-    await start_and_await_workflow(client, input_data, workflow_id)
-
     assert_stop(postgres_engine, model_id)
-    
+    mock_mlflow_models['transform_model'].predict.assert_not_called()
-    print("\n[TEST] ✓ All assertions passed!")
 
 
 @pytest.mark.asyncio
@@ -255,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,
 ):
-    """
+    """Input gate REPEAT with existing history."""
-    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
-    """
     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'
+    input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger REPEAT...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-input-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")
-        )
     )
+    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)
+
-    
+@pytest.mark.asyncio
-    return model
+@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
@@ -317,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'
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at transform gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-continue')
-    workflow_id = f'test-transform-continue-{datetime.now().timestamp()}'
+    )
-    
-    await start_and_await_workflow(client, input_data, workflow_id)
     assert_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!")
 
 
 @pytest.mark.asyncio
@@ -369,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'
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'STOP'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger STOP at transform gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-stop')
-    workflow_id = f'test-transform-stop-{datetime.now().timestamp()}'
+    )
-    
-    await start_and_await_workflow(client, input_data, workflow_id)
     assert_stop(postgres_engine, model_id)
-
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
-    print("\n[TEST] ✓ All assertions passed!")
-
 
 
 @pytest.mark.asyncio
@@ -418,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'
+    input_data['mlflow_transform_filters']['API_ERROR']['POLICY'] = 'REPEAT'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger REPEAT at transform gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-transform-repeat')
-    workflow_id = f'test-transform-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.mark.asyncio
-@pytest.fixture
+@pytest.mark.integration
-def bad_predict_model(
+async def test_scenario_2_2_4_transform_content_gate_nan_values_stop(
-    patch_mlflow,
+    temporal_test_env: WorkflowEnvironment,
-    mock_mlflow_models
+    temporal_worker: Worker,
+    test_activities: Activities,
+    postgres_engine,
+    mock_mlflow_models,
 ):
-    model = MagicMock(
+    """mlflow_content_gate triggers STOP when transform output is all NaN (NAN_VALUES filter)."""
-        predict=MagicMock(
+    client = temporal_test_env.client
-            side_effect=Exception("Bad predict model")
+    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')
     )
-
+    assert_stop(postgres_engine, model_id)
-    def mock_sklearn_load_model(model_uri):
+    mock_mlflow_models['predict_model'].predict.assert_not_called()
-        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
@@ -490,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'
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'CONTINUE'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger CONTINUE at predict gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-continue')
-    workflow_id = f'test-predict-continue-{datetime.now().timestamp()}'
+    )
-    
-    await start_and_await_workflow(client, input_data, workflow_id)
     assert_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!")
 
 
 @pytest.mark.asyncio
@@ -541,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'
+    input_data['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'STOP'
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should stop at predict gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-stop')
-    workflow_id = f'test-predict-stop-{datetime.now().timestamp()}'
+    )
-    
-    await start_and_await_workflow(client, input_data, workflow_id)
     assert_stop(postgres_engine, model_id)
 
-    print("\n[TEST] ✓ All assertions passed!")
-
 
 @pytest.mark.asyncio
 @pytest.mark.integration
@@ -587,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['mlflow_predict_filters']['API_ERROR']['POLICY'] = 'REPEAT'
-    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']  # REPEAT first
+    input_data['path_priority'] = ['REPEAT', 'STOP', 'CONTINUE']
-
+    await start_and_await_workflow(
-    print("\n[TEST] 2. Starting workflow that should trigger REPEAT at predict gate...")
+        client, PredictionsBatch.run, input_data, make_workflow_id('test-predict-repeat')
-    workflow_id = f'test-predict-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.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)

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
@@ -73,6 +74,17 @@ class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
         """
         metrics_controller = MetricsController(logger=logger)
 
+        minio_repository = MinioRepository(
+            endpoint=minio_config['endpoint_url'],
+            access_key=minio_config['access_key'],
+            secret_key=minio_config['secret_key'],
+            bucket=minio_config['default_bucket'],
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+            secure=minio_config['secure'],
+        )
+
         # Initialize parent classes
         Storage.__init__(
             self,
@@ -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,

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
+        PI Web API metrics in laborious use the shared ``CORE_LABELS`` from
-        a convenience overload behavior:
+        ``sientia_do``, which includes ``operation_type``. For this reason,
-         - When operation_type is provided, it behaves exactly like the base class,
+        operation_type must always be present in emitted labels.
-           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.
 
         Args:
-         - metadata (dict[str, Any]): Workflow execution metadata used to derive labels
+         - metadata (dict[str, Any]): Workflow execution metadata used to derive labels.
-         - operation_type (str | None): Optional operation type label. If None, the
+         - operation_type (str): Operation type label for metric cardinality.
-           operation_type key will be removed from the returned labels.
 
         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:
         """
@@ -256,8 +248,13 @@ class API(SientiaMonitoring):
                 metadata=metadata,
             )
 
-            data['prediction_confidence'] = confidence
+            # Preserve incoming confidence/comments on successful PI writes.
-            data['comments'] = message
+            # 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()
@@ -282,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,
             )

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
@@ -10,16 +13,16 @@ 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.dataframe_debug import build_dataframe_debug_message
     from laborious.utils.filters.conditional_filters import (
         filter_empty_data,
         filter_specific_variables_null_values,
     )
     from laborious.utils.filters.mlflow_filters import api_error_filter, nan_values_filter
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
 
 # Strongly-typed filter function signatures
 InputFilterFunc = Callable[[DataFrame, dict[str, Any]], bool]
@@ -63,7 +66,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 +86,15 @@ class Gates(SientiaMonitoring):
         mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
     """
 
+    minio_repository: MinioRepository | None = None
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
     def __init__(
         self,
-        logger: Logger,
+        minio_repository: MinioRepository | None = None,
-        notification_handler: NotificationHandler,
+        logger: Logger | None = None,
-        metrics_controller: MetricsController,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
     ):
         """
         Initialize data quality gates with logging and notification capabilities.
@@ -99,17 +106,54 @@ 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()
 
+    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
+        """
+        Log dataframe content only when row count is below the configured threshold
+
+        Args:
+         - message (str): Base log message to identify the dataframe in logs
+         - data (Any): Dataframe-like payload to be logged
+         - metadata (dict[str, Any]): Workflow metadata for contextual logging
+        """
+        self.debug(
+            build_dataframe_debug_message(
+                message=message,
+                data=data,
+                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
+            ),
+            metadata,
+        )
+
+    @staticmethod
+    def _read_filter_entry(config: dict[str, Any]) -> tuple[str, dict[str, Any]]:
+        """
+        Read filter policy/config keys in a case-insensitive way.
+
+        Args:
+            config (dict[str, Any]): Filter configuration dictionary.
+
+        Return:
+            tuple[str, dict[str, Any]]: Parsed policy and config payload.
+        """
+        normalized = {str(key).upper(): value for key, value in config.items()}
+        policy = normalized['POLICY']
+        filter_config = normalized.get('CONFIG', {})
+        return policy, filter_config
+
     @activity.defn(name='input_gate')
     async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
         """
@@ -149,12 +193,13 @@ class Gates(SientiaMonitoring):
         self.info('Performing input gate...', metadata)
 
         filters = input_data['filters']
-        data = DataFrame(input_data['data'])
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(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
@@ -162,10 +207,11 @@ class Gates(SientiaMonitoring):
             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(
@@ -183,6 +229,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')
@@ -221,32 +270,41 @@ class Gates(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Performing mlflow response gate...', metadata)
-
+        raw_data = input_data['data']
         filters = input_data['filters']
-        data = input_data['data']
+
+        self.debug(
+            f'Input data: \n {create_sample_dict(raw_data, max_items=5, max_depth=5)}', metadata
+        )
+        self.debug(f'Filters: {filters}', metadata)
+
+        payload = MinioDataFramePayload.from_dict(raw_data)
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         gate_type = input_data['type']
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(f'Input data: \n {create_sample_dict(data, max_items=5, max_depth=5)}', metadata)
-        self.debug(f'Filters: {filters}', metadata)
-
         comments = []
+
+        status = payload.status or {}
+
         for fil, config in filters.items():
             if fil not in mlflow_response_filter_functions:
                 continue
+            policy, filter_config = self._read_filter_entry(config)
             try:
-                if mlflow_response_filter_functions[fil](data, config):
+                if mlflow_response_filter_functions[fil](status, filter_config):
-                    filter_output.append(config['policy'])
+                    filter_output.append(policy)
-                    comments.append(data['content']['message'])
+                    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()
@@ -265,6 +323,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,21 +366,25 @@ class Gates(SientiaMonitoring):
         self.info('Performing mlflow content gate...', metadata)
 
         filters = input_data['filters']
-        data = DataFrame(input_data['data'])
+
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         gate_type = input_data['type']
         path_priority = input_data['path_priority']
 
         filter_output = []
 
-        self.debug(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():
             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):
+                if mlflow_content_filter_functions[fil](data, filter_config):
-                    filter_output.append(config['policy'])
+                    filter_output.append(policy)
                     await self.send_notification_async(
                         metadata=metadata,
                         notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
@@ -349,6 +414,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 +471,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 +506,8 @@ class Gates(SientiaMonitoring):
 
         self.info('Formatting transformed data...', metadata)
 
-        data = DataFrame(input_data['data'])
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
 
         data['timestamp'] = data.index
         data = data.reset_index(drop=True)
@@ -446,7 +515,15 @@ 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,
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
 
     @activity.defn(name='format_prediction')
     async def format_prediction(self, input_data: dict[str, Any]) -> dict:
@@ -458,6 +535,8 @@ class Gates(SientiaMonitoring):
         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
@@ -465,7 +544,7 @@ class Gates(SientiaMonitoring):
         Args:
             input_data (dict): Input data containing:
                 - data (dict[str, Any]): Raw prediction data to format
-                - timestamp (str): Default timestamp if data lacks timestamp column
+                - timestamp (str): Timestamp of the data
                 - model_id (str): Unique identifier for the ML model
                 - prediction_confidence (float): Confidence score for the prediction
                 - prediction_store_policy (str): Storage policy in format 'type:value'
@@ -474,17 +553,19 @@ class Gates(SientiaMonitoring):
             dict: Formatted prediction data ready for storage and export
         """
         metadata = input_data['metadata']
+        last_timestamp = input_data['timestamp']
         prediction_store_policy = input_data['prediction_store_policy']
         self.info('Formatting prediction...', metadata)
 
-        data = DataFrame(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
         data['timestamp'] = data.index
         data = data.reset_index(drop=True)
 
         self.debug(f'Prediction store policy: {prediction_store_policy}', metadata)
-        self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
+        self._debug_dataframe('Prediction data:', data, metadata)
 
         policy_type, policy_value = self.get_prediction_store_policy(
             prediction_store_policy, metadata
@@ -507,7 +588,11 @@ class Gates(SientiaMonitoring):
             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']
@@ -517,7 +602,7 @@ class Gates(SientiaMonitoring):
         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()
 
@@ -632,50 +717,10 @@ class Gates(SientiaMonitoring):
             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()
 
-    @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]):
         """
@@ -707,34 +752,29 @@ class Gates(SientiaMonitoring):
 
         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={
+            tags=core_tags,
-                'pod_id': self.pod_id,
-                'model_name': metadata['model_name'],
-                'workflow_name': metadata['workflow_name'],
-            },
         )
 
         await self.emit_metric(
             metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
             method='set',
-            tags={
+            tags=core_tags,
-                'pod_id': self.pod_id,
-                'model_name': metadata['model_name'],
-                'workflow_name': metadata['workflow_name'],
-            },
             value=prediction_confidence,
         )
 
         await self.emit_metric(
             metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
             method='observe',
-            tags={
+            tags=core_tags,
-                'pod_id': self.pod_id,
-                'model_name': metadata['model_name'],
-                'workflow_name': metadata['workflow_name'],
-            },
             value=response_time,
         )
 
@@ -745,9 +785,7 @@ class Gates(SientiaMonitoring):
                         metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
                         method='observe',
                         tags={
-                            'pod_id': self.pod_id,
+                            **core_tags,
-                            'model_name': metadata['model_name'],
-                            'workflow_name': metadata['workflow_name'],
                             'opc_server_id': server_id,
                             'tag': tag,
                         },
@@ -757,9 +795,7 @@ class Gates(SientiaMonitoring):
                     await self.emit_metric(
                         metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
                         tags={
-                            'pod_id': self.pod_id,
+                            **core_tags,
-                            'model_name': metadata['model_name'],
-                            'workflow_name': metadata['workflow_name'],
                             'opc_server_id': server_id,
                             'tag': tag,
                         },

laborious/activities/mlflow.py

@@ -5,12 +5,12 @@ with workflow.unsafe.imports_passed_through():
     from typing import Any
 
     import numpy as np
-    from pandas import DataFrame, 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,
@@ -19,11 +19,13 @@ 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.dataframe_debug import build_dataframe_debug_message
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
+    from laborious.utils.repository.minio_manager import MinioManager
     from laborious.utils.repository.model_repository import MLFlowRepository
 
 
-class MLFlow(SientiaMonitoring):
+class MLFlow(MinioManager):
     """
     MLFlow integration activities for model inference operations.
 
@@ -42,16 +44,18 @@ class MLFlow(SientiaMonitoring):
         model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
     """
 
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
     def __init__(
         self,
         mlflow_host: str,
         mlflow_port: int,
         mlflow_username: str,
-        minio_config: dict[str, Any],
         mlflow_password: str,
-        logger: Logger,
+        minio_repository: MinioRepository | None = None,
-        notification_handler: NotificationHandler,
+        logger: Logger | None = None,
-        metrics_controller: MetricsController,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
     ):
         """
         Initialize MLFlow activities with server configuration.
@@ -67,7 +71,9 @@ 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 +88,35 @@ 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()
 
+    def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
+        """
+        Log dataframe content only when row count is below the configured threshold
+
+        Args:
+         - message (str): Base log message to identify the dataframe in logs
+         - data (Any): Dataframe-like object expected to expose shape and to_csv
+         - metadata (dict[str, Any]): Workflow metadata for contextual logging
+        """
+        self.debug(
+            build_dataframe_debug_message(
+                message=message,
+                data=data,
+                max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
+            ),
+            metadata,
+        )
+
     @activity.defn(name='request_transform')
-    async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
+    async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
         """
         Transform input data using MLFlow models.
 
@@ -138,12 +147,14 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Transforming data...', metadata)
-        data = DataFrame(input_data['data'])
+
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug('Raw input data:', metadata)
+        self._debug_dataframe('Raw input data:', data, metadata)
-        self.debug(data.head(5).to_string(), metadata)
 
         # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
         data = data.sort_values('created_at', ascending=False).drop_duplicates(
@@ -157,7 +168,9 @@ class MLFlow(SientiaMonitoring):
         data.columns.name = None
         data.index.name = None
 
-        self.debug(f'Processed input data: \n {data.to_csv()}', metadata)
+        data['timestamp'] = data.index
+
+        self._debug_dataframe('Processed input data:', data, metadata)
 
         # Request transformation from MLFlow model
         response_data = await self.model_monitoring_repository.transform(
@@ -176,10 +189,33 @@ 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,
+                last_timestamp=payload.last_timestamp,
+                logger=self.logger,
+            )
+
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe=response_data['content'],
+            minio_repo=self.minio_repository,
+            model_name=model_name,
+            operation='transform',
+            workflow_metadata=metadata,
+            status={
+                'success': True,
+            },
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
 
     @activity.defn(name='request_predict')
-    async def request_predict(self, input_data: dict[str, Any]) -> dict[str, Any]:
+    async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
         """
         Execute predictions using MLFlow models.
 
@@ -210,11 +246,14 @@ class MLFlow(SientiaMonitoring):
         """
         metadata = input_data['metadata']
         self.info('Predicting data...', metadata)
-        data = DataFrame(input_data['data'])
+
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
         model_name = input_data['model_name']
         model_config = input_data.get('model_config', {})
 
-        self.debug(f'Input data for: \n {data.head(5).to_string()}', metadata)
+        self._debug_dataframe('Input data for prediction:', data, metadata)
 
         # Convert numpy.nan to None for model compatibility
         data.replace(np.nan, None, inplace=True)
@@ -236,7 +275,30 @@ 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,
+                last_timestamp=payload.last_timestamp,
+                logger=self.logger,
+            )
+
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe=response_data['content'],
+            minio_repo=self.minio_repository,
+            model_name=model_name,
+            operation='predict',
+            workflow_metadata=metadata,
+            status={
+                'success': True,
+            },
+            last_timestamp=payload.last_timestamp,
+            logger=self.logger,
+        )
 
     @activity.defn(name='retrain_model')
     async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
@@ -275,14 +337,12 @@ 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(
+            # Payload-based retrain input (inline dict or MinIO offloaded).
-                object_key=object_key, metadata=metadata
+            payload = MinioDataFramePayload.from_dict(input_data['data'])
-            )
+            data = await payload.retrieve(self.minio_repository, metadata)
+
         except Exception as e:
             trace = traceback.format_exc()
             await self.send_notification_async(

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(
@@ -31,6 +32,8 @@ class ModelMetrics(SientiaMonitoring):
     This class provides activities for writing metrics to the Prometheus monitoring system.
     """
 
+    _MAX_DEBUG_DATAFRAME_ROWS = 100
+
     def __init__(
         self,
         logger: Logger,
@@ -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,11 @@ class ModelMetrics(SientiaMonitoring):
 
         model_analysis = ModelAnalysis(config=config)
 
-        self.debug(
+        self._debug_dataframe(
-            f'Reference data: Size {reference_data.shape} \n{reference_data.head(5).to_string()}',
+            f'Reference data: Size {reference_data.shape}', reference_data, metadata
-            metadata,
         )
 
-        self.debug(
+        self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata)
-            f'Target data: Size {target_data.shape} \n{target_data.head(5).to_string()}', metadata
-        )
 
         core_labels = self.get_core_labels(metadata, operation_type='detect_univariate_drift')
         start_time = time.time()
@@ -142,9 +160,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(
+        self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
-            f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata
-        )
 
         return drift_df
 
@@ -274,11 +290,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(
+        self._debug_dataframe(f'Drift dataframe: Size {drift_df.shape}', drift_df, metadata)
-            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)
 
         return drift_df.to_dict(orient='records')
 
@@ -347,8 +359,6 @@ class ModelMetrics(SientiaMonitoring):
         data['data_size'] = data_size
         data['interval_minutes'] = interval_minutes
 
-        self.debug(
+        self._debug_dataframe(f'Simple metrics dataframe: Size {data.shape}', data, metadata)
-            f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata
-        )
 
         return data.to_dict(orient='records')

laborious/activities/storage.py

@@ -1,8 +1,11 @@
 from temporalio import activity, workflow
 
+from laborious.utils.repository.minio_manager import MinioManager
+
 with workflow.unsafe.imports_passed_through():
     # Extend the Temporal Postgres activities for convenient query -> MinIO export
     import traceback
+    from datetime import timedelta
     from typing import Any
 
     import pandas as pd
@@ -10,18 +13,23 @@ 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 sientia_do.temporal.constants import now
 
-    from laborious.utils.repository.minio_repository import MinioRepository
+    from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
+
+_LOAD_QUERY_OFFLOAD_SKIP_KEYS = frozenset({'model_name', 'key_prefix', 'size_threshold_bytes'})
 
 
-class Storage(Postgres):
+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 +39,15 @@ class Storage(Postgres):
         dbname: str,
         min_connections: int,
         max_connections: int,
-        minio_config: dict[str, Any],
+        retention_hours: int = 24,
-        logger: Logger,
+        minio_repository: MinioRepository | None = None,
-        notification_handler: NotificationHandler,
+        logger: Logger | None = None,
-        metrics_controller: MetricsController,
+        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,90 +60,150 @@ class Storage(Postgres):
             metrics_controller=metrics_controller,
         )
 
-        if not hasattr(self, 'minio_repository'):
+        MinioManager.__init__(
-            self.minio_repository: MinioRepository | None = None
+            self, minio_repository, logger, notification_handler, metrics_controller
+        )
 
-        if self.minio_repository is None:
+    @activity.defn(name='load_query_with_minio_offload')
-            self.minio_repository = MinioRepository(
+    async def load_query_with_minio_offload(
-                logger=logger,
+        self, input_data: dict[str, Any]
-                notification_handler=notification_handler,
+    ) -> MinioDataFramePayload:
-                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,
-            )
-
-    @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.
+        Run the custom SQL load, then return a MinIO-aware dataframe wire dict.
 
         Args (input_data):
-            metadata (dict): Workflow metadata
+            metadata (dict): Workflow metadata (same as load_custom_query).
-            query (str): SQL query
+            query (str): SQL query.
-            model_name (str): Model name for object naming
+            datetime_columns (list[str], optional): Datetime column names.
-            object_prefix (str, optional): Prefix inside bucket (default: datasets/retrain)
+            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: { success: bool, object_name: str, uri: str }
+            dict[str, Any]: Flat ``MinioDataFramePayload`` dict or ``success: False`` on failure.
         """
+        if self.minio_repository is None:
+            raise ValueError('Minio repository not initialized')
 
+        metadata: dict = input_data.get('metadata', {})
+        model_name = input_data['model_name']
+
+        rows = await self.load_custom_query(
+            input_data,
+        )
+        if not rows:
+            self.error(
+                'load_query_with_minio_offload failed: No data returned from query', metadata
+            )
+            dataframe = None
+        else:
+            dataframe = pd.DataFrame(rows)
+
+        return await MinioDataFramePayload.from_dataframe(
+            dataframe,
+            minio_repo=self.minio_repository,
+            workflow_metadata=metadata,
+            model_name=model_name,
+            operation='initial',
+            logger=self.logger,
+        )
+
+    @activity.defn(name='export_payload_to_postgres')
+    async def export_payload_to_postgres(self, input_data: dict[str, Any]) -> dict:
+        """
+        Export a payload to PostgreSQL.
+        """
+        metadata = input_data.get('metadata')
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
+        data = await payload.retrieve(self.minio_repository, metadata)
+
+        return await self.export_data_to_postgres(
+            {
+                **input_data,
+                'data': data,
+            }
+        )
+
+    @activity.defn(name='cleanup_minio_objects_expired')
+    async def cleanup_minio_objects_expired(self, input_data: dict[str, Any]) -> dict[str, Any]:
+        """
+        Delete objects under the given prefixes that are older than the retention window.
+
+        Args (input_data):
+            metadata (dict): Workflow metadata for logging and metrics.
+            prefixes (list[str]): Key prefixes to scan (one level or subtree per prefix).
+
+        Returns:
+            dict[str, Any]: ``success``, ``deleted_count``, and optional ``message``.
+        """
         if self.minio_repository is None:
             raise ValueError('Minio repository not initialized')
 
         metadata = input_data.get('metadata', {})
-        object_prefix = input_data.get('object_prefix', 'datasets/retrain')
+        payload = MinioDataFramePayload.from_dict(input_data['data'])
-
+        prefix = payload.cleanup_prefix()
-        timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
+        base = now()
-        object_name = f'{object_prefix}_{timestamp}.parquet'
+        cutoff = (base.replace(tzinfo=None) if base.tzinfo else base) - timedelta(
-        uri = f's3://{self.minio_repository.minio_bucket}/{object_name}'
+            hours=self.retention_hours
+        )
 
+        report: dict[str, Any] = {
+            'failed': {},
+            'deleted': {},
+            'failed_count': 0,
+            'deleted_count': 0,
+        }
         try:
-            data = await self.load_custom_query(input_data)
+            keys = await self.minio_repository.list_objects(
-            if not data:
+                prefix=prefix,
-                self.error('query_to_minio failed: No data returned from query', metadata)
+                recursive=True,
-                return {'success': False, 'message': 'No data returned from query'}
+                metadata=metadata,
-
-            # 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
             )
-
+            for key in keys:
-            return {'success': True, 'object_key': object_name, 'uri': uri}
+                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_STORING_QUERY_TO_MINIO',
+                notification_id='ERROR_CLEANUP_MINIO_OBJECTS_EXPIRED',
-                message=f'Error storing query to MinIO: {e}',
+                message=f'Error cleaning up MinIO objects: {e}',
-                block='query_to_minio',
+                block='cleanup_minio_objects_expired',
                 level=NotificationLevel.ERROR,
                 attachment_content=trace,
             )
-
             self.error(trace, metadata)
+        else:
+            # Cleanup success is expected in normal flow; avoid noisy INFO notifications
+            # that do not impact behavior and can flood observability in test runs.
+            self.info('MinIO objects cleaned up successfully', metadata)
 
-            return {'success': False, 'message': str(e)}
+        return report
 
     def close(self) -> None:
         """Close Storage resources (MinIO client and Postgres engine)."""
-        try:
+        Postgres.close(self)
-            if hasattr(self, 'minio_repository') and self.minio_repository is not None:
+        MinioManager.close(self)
-                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')
 
     def __del__(self):
         self.close()

laborious/metrics.py

@@ -18,15 +18,14 @@ 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
 """
 
 from prometheus_client import Counter, Gauge, Histogram
-from sientia_do.observability.metrics import (
+from sientia_do.observability.metrics import CORE_LABELS
-    CORE_LABELS as SIENTIA_CORE_LABELS,
-)
 
 # Application health metric
 APP_UP = Gauge(
@@ -35,9 +34,6 @@ APP_UP = Gauge(
     ['pod_id'],
 )
 
-# Core labels used across multiple metrics
-CORE_LABELS = ['pod_id', 'model_name', 'workflow_name']
-
 # Prediction operation metrics
 PREDICTIONS_WRITTEN_COUNT = Counter(
     'laborious_predictions_written_count',
@@ -60,45 +56,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 ==================
 
@@ -137,58 +94,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,
 )
 
 

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
     """
@@ -86,6 +86,7 @@ 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',
     }

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}'
+    )

laborious/utils/models/minio_dataframe_payload.py

@@ -0,0 +1,348 @@
+"""
+MinIO-backed DataFrame payload for Temporal workflows.
+
+Data is never stored as a pandas ``DataFrame`` field on the dataclass.
+Instead, the DataFrame is only provided as an input to:
+`from_dataframe` / `from_dataframe_to_dict`.
+
+At build time, the DataFrame is evaluated for its serialized size; if it exceeds
+the configured threshold, it is serialized to parquet bytes and uploaded to MinIO.
+Otherwise, it is inlined as a Temporal-friendly ``dict``.
+"""
+
+import pickle
+import re
+from collections.abc import Hashable
+from dataclasses import dataclass
+from datetime import datetime
+from io import BytesIO
+from os import getenv
+from typing import Any, Literal
+
+from pandas import DataFrame, read_parquet
+from sientia_do.observability.logger import Logger
+from sientia_do.repository.minio_repository import MinioRepository
+from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, DATETIME_FORMAT_WITH_TZ, now
+
+# Keys that are part of the serialized wire format (not arbitrary metadata).
+_SERIALIZED_FIELD_KEYS = frozenset({'data', 'bucket', 'object_key', 'object_prefix', 'uri'})
+
+_OBJECT_TIMESTAMP_PATTERN = re.compile(
+    r'-(?:initial|transform)-(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})\.parquet$'
+)
+
+OFFLOAD_THRESHOLD_BYTES = int(
+    float(getenv('SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES', '1.5')) * 1024 * 1024
+)
+
+# Relative prefix used for storing offloaded prediction datasets in MinIO.
+# It is also the root directory for retention cleanup listing.
+PREDICTION_DATASETS_PREFIX = 'prediction_datasets'
+
+OperationKind = Literal['initial', 'transform', 'predict']
+
+
+def _build_object_key(
+    model_name: str, operation: OperationKind, timestamp: str
+) -> tuple[str, str | None]:
+    """
+    Build the MinIO object key and the directory prefix used for retention listing.
+
+    Args:
+        model_name: Registered model name used in the pipeline.
+        operation: Either initial (pre-transform load) or transform (post-MLFlow transform).
+        timestamp: Filename timestamp segment from DATETIME_FORMAT_FILENAME.
+
+    Return:
+        tuple[str, str | None]: Full object key and normalized prefix (or None if at bucket root).
+    """
+    # Naming convention:
+    #   - Directory is always `prediction_datasets/<model_name>`
+    #   - Filename follows the retention-parsing pattern
+    basename = f'{model_name}-{operation}-{timestamp}.parquet'
+    model_dir = model_name.strip().strip('/')
+    prefix = f'{PREDICTION_DATASETS_PREFIX}/{model_dir}'
+    return f'{prefix}/{basename}', prefix
+
+
+@dataclass
+class MinioDataFramePayload:
+    """
+    Serializable payload after a DataFrame was evaluated: inline tabular dict and/or MinIO keys.
+
+    Build from a live DataFrame only via `from_dataframe` / `from_dataframe_to_dict`.
+    Rehydrate from Temporal via `from_dict`. The DataFrame is not a field on this class.
+    """
+
+    last_timestamp: str
+    status: dict[str, Any] | None = None
+    data: dict[Hashable, Any] | None = None
+    bucket: str | None = None
+    object_key: str | None = None
+    object_prefix: str | None = None
+    uri: str | None = None
+
+    @staticmethod
+    def _debug(
+        logger: Logger | None,
+        message: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        """
+        Emit debug logs only when logger is provided
+
+        Args:
+         - logger (Logger | None): Logger instance used for debug messages
+         - message (str): Message to be logged
+         - metadata (dict[str, Any] | None): Optional workflow metadata context
+        """
+        if logger is None:
+            return
+        logger.custom_debug(message, metadata)
+
+    @classmethod
+    def from_dict(cls, raw: 'dict[str, Any] | MinioDataFramePayload') -> 'MinioDataFramePayload':
+        """
+        Reconstruct a MinioDataFramePayload from a plain dict produced by Temporal serialization.
+
+        Temporal converts dataclass return values into plain dicts when crossing
+        workflow/activity boundaries. This method rebuilds the typed instance so
+        that methods like ``retrieve``, ``cleanup_prefix`` and ``has_data`` are
+        available on the receiving side.
+
+        If the argument is already a MinioDataFramePayload, it is returned as-is.
+
+        Args:
+            raw: Dict with keys matching the dataclass fields
+                (last_timestamp, status, data, bucket, object_key, object_prefix, uri),
+                or an existing MinioDataFramePayload instance.
+
+        Return:
+            MinioDataFramePayload: Reconstructed (or original) instance.
+        """
+        if isinstance(raw, MinioDataFramePayload):
+            return raw
+        return cls(
+            last_timestamp=raw['last_timestamp'],
+            status=raw.get('status'),
+            data=raw.get('data'),
+            bucket=raw.get('bucket'),
+            object_key=raw.get('object_key'),
+            object_prefix=raw.get('object_prefix'),
+            uri=raw.get('uri'),
+        )
+
+    @staticmethod
+    def estimate_size_bytes(
+        df: DataFrame,
+        metadata: dict[str, Any] | None = None,
+        logger: Logger | None = None,
+    ) -> int:
+        """
+        Approximate serialized size of the DataFrame as the default-orient dict.
+
+        Args:
+            df: DataFrame whose tabular content size is estimated.
+
+        Return:
+            int: Estimated size in bytes (pickle of dict representation).
+        """
+        try:
+            size = len(pickle.dumps(df.to_dict()))
+        except Exception:
+            size = len(pickle.dumps(df))
+
+        MinioDataFramePayload._debug(
+            logger,
+            f'DataFrame size: {size} bytes',
+            metadata,
+        )
+        return size
+
+    @staticmethod
+    def parse_object_timestamp(object_key: str) -> datetime | None:
+        """
+        Parse the timestamp embedded in the object key basename (before .parquet).
+
+        Args:
+            object_key: S3/MinIO object key whose basename follows
+                ``{model}-{initial|transform}-{DATETIME_FORMAT_FILENAME}.parquet``.
+
+        Return:
+            datetime | None: Parsed UTC-naive datetime from the key, or None if not matched.
+        """
+        basename = object_key.rsplit('/', 1)[-1]
+        match = _OBJECT_TIMESTAMP_PATTERN.search(basename)
+        if not match:
+            return None
+        try:
+            return datetime.strptime(match.group(1), DATETIME_FORMAT_FILENAME)
+        except ValueError:
+            return None
+
+    def cleanup_prefix(self) -> str | None:
+        """
+        Return True if cleanup is enabled for this payload.
+        """
+        if self.object_key is not None and self.data is None:
+            return self.object_prefix
+        return None
+
+    def has_data(self) -> bool:
+        """
+        Return True if the payload has some data internally or in MinIO.
+        """
+        return (self.data is not None and self.data != {}) or self.object_key is not None
+
+    @classmethod
+    async def from_dataframe(
+        cls,
+        dataframe: DataFrame | None,
+        minio_repo: MinioRepository,
+        model_name: str,
+        operation: OperationKind,
+        status: dict[str, Any] | None = None,
+        workflow_metadata: dict | None = None,
+        last_timestamp: str | None = None,
+        logger: Logger | None = None,
+    ) -> 'MinioDataFramePayload':
+        """
+        Evaluate the DataFrame size, then either inline dict or upload parquet to MinIO.
+
+        The DataFrame is not stored on the returned instance.
+
+        Args:
+            dataframe: Tabular data to evaluate and persist (inline or MinIO).
+            metadata: Small metadata dict merged into the payload (e.g. success, message).
+            minio_repo: sientia_do MinioRepository (or compatible) with `upload_file()`.
+            workflow_metadata: Metadata passed to MinIO store for logging/metrics.
+            model_name: Registered model name used in the object basename.
+            operation: Either ``initial`` (query load) or ``transform`` (post-transform).
+            key_prefix: Backward-compatible parameter (currently ignored for object naming).
+            size_threshold_bytes: Byte limit before offload. When None, the module-level
+                environment-derived default is used.
+
+        Return:
+            MinioDataFramePayload: Instance with data and/or MinIO fields set.
+        """
+
+        if dataframe is None or dataframe.empty:
+            cls._debug(
+                logger,
+                'MinioDataFramePayload.from_dataframe received empty dataframe, returning empty payload',
+                workflow_metadata,
+            )
+            return cls(
+                data=None, last_timestamp=now().strftime(DATETIME_FORMAT_WITH_TZ), status=status
+            )
+
+        if last_timestamp is None:
+            last_timestamp = max(dataframe['timestamp'].values.tolist())
+
+        dataframe_size = cls.estimate_size_bytes(dataframe, workflow_metadata, logger)
+        cls._debug(
+            logger,
+            (
+                f'MinioDataFramePayload.from_dataframe estimated size: {dataframe_size} bytes '
+                f'(threshold: {OFFLOAD_THRESHOLD_BYTES} bytes)'
+            ),
+            workflow_metadata,
+        )
+
+        if dataframe_size <= OFFLOAD_THRESHOLD_BYTES:
+            cls._debug(
+                logger,
+                'MinioDataFramePayload.from_dataframe using inline payload',
+                workflow_metadata,
+            )
+            return cls(data=dataframe.to_dict(), last_timestamp=last_timestamp, status=status)
+
+        timestamp = now().strftime(DATETIME_FORMAT_FILENAME)
+        object_key, object_prefix = _build_object_key(model_name, operation, timestamp)
+        cls._debug(
+            logger,
+            (
+                'MinioDataFramePayload.from_dataframe offloading payload to MinIO '
+                f'with key {object_key}'
+            ),
+            workflow_metadata,
+        )
+
+        # Upload using the relative object key. The upstream repository will
+        # prefix it internally under its MinIO namespace.
+        parquet_buffer = BytesIO()
+        dataframe.to_parquet(parquet_buffer, engine='pyarrow', index=True)
+        file_bytes = parquet_buffer.getvalue()
+
+        upload_result = await minio_repo.upload_file(
+            file_bytes=file_bytes,
+            relative_key=object_key,
+            metadata=workflow_metadata,
+        )
+
+        bucket = minio_repo.bucket
+        object_key_full = upload_result.get('minio_object_name', object_key)
+        uri = f's3://{bucket}/{object_key_full}' if bucket else None
+        cls._debug(
+            logger,
+            f'MinioDataFramePayload.from_dataframe upload completed: {uri}',
+            workflow_metadata,
+        )
+
+        return cls(
+            data=None,
+            bucket=bucket,
+            object_key=object_key_full,
+            object_prefix=object_prefix,
+            uri=uri,
+            last_timestamp=last_timestamp,
+            status=status,
+        )
+
+    async def retrieve(
+        self,
+        minio_repo: MinioRepository,
+        workflow_metadata: dict[str, Any] | None = None,
+        logger: Logger | None = None,
+    ) -> DataFrame:
+        """
+        Load parquet from MinIO when object_key is set and populate inline data.
+
+        Args:
+            minio_repo: sientia_do MinioRepository (or compatible) with download_file().
+            workflow_metadata: Metadata passed to MinIO read for logging/metrics.
+
+        Return:
+            dict[str, Any]: Flat dict with data filled (same keys as to_dict after load).
+        """
+        if self.data is not None:
+            self._debug(
+                logger,
+                'MinioDataFramePayload.retrieve using inline payload data',
+                workflow_metadata,
+            )
+            return DataFrame(self.data)
+
+        if not self.has_data():
+            self._debug(
+                logger,
+                'MinioDataFramePayload.retrieve found no payload data, returning empty dataframe',
+                workflow_metadata,
+            )
+            return DataFrame()
+
+        self._debug(
+            logger,
+            f'MinioDataFramePayload.retrieve downloading object from MinIO: {self.object_key}',
+            workflow_metadata,
+        )
+        file_bytes = await minio_repo.download_file(
+            object_name=self.object_key, metadata=workflow_metadata
+        )
+        df = read_parquet(BytesIO(file_bytes))
+        self._debug(
+            logger,
+            f'MinioDataFramePayload.retrieve loaded dataframe from MinIO with shape {df.shape}',
+            workflow_metadata,
+        )
+        return df

laborious/utils/repository/minio_manager.py

@@ -0,0 +1,32 @@
+from sientia_do.notifications.handlers import NotificationHandler
+from sientia_do.observability.logger import Logger
+from sientia_do.observability.metrics_controller import MetricsController
+from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+from sientia_do.repository.minio_repository import MinioRepository
+
+
+class MinioManager(SientiaMonitoring):
+    minio_repository: MinioRepository | None = None
+
+    def __init__(
+        self,
+        minio_repository: MinioRepository | None = None,
+        logger: Logger | None = None,
+        notification_handler: NotificationHandler | None = None,
+        metrics_controller: MetricsController | None = None,
+    ):
+        if self.minio_repository is None:
+            self.minio_repository = minio_repository
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
+
+    def close(self) -> None:
+        """
+        Close the MinioManager and clean up resources.
+        """
+        if self.minio_repository is not None:
+            try:
+                self.minio_repository.close()
+            finally:
+                self.minio_repository = None
+
+        SientiaMonitoring.shutdown(self)

laborious/utils/repository/minio_repository.py

@@ -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)

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
     """
@@ -452,6 +473,10 @@ 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
+            )
         else:
             if model_type == 'predict':
                 model = await self.load_predict_model(model_name, metadata, flavor)
@@ -1134,7 +1159,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.
 
@@ -1168,7 +1193,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 +1211,8 @@ class MLFlowRepository(SientiaMonitoring):
                 metadata=metadata,
             )
 
-            self.debug(
+            self._debug_dataframe(
-                f'Data received from model transformation: {transformed_data.head(5).to_csv()}',
+                'Data received from model transformation:', transformed_data, metadata
-                metadata,
             )
 
             # transformed_data.to_csv(
@@ -1196,7 +1220,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 {
@@ -1254,9 +1278,7 @@ class MLFlowRepository(SientiaMonitoring):
             input_index = data.index
             start_time = datetime.now()
 
-            self.debug(
+            self._debug_dataframe('Data received for model prediction:', data, metadata)
-                f'Data received for model prediction: {data.to_dict(orient="records")}', metadata
-            )
 
             # data.to_csv(
             #     f"tmp/treated_data_{model_name}.csv", index=True)
@@ -1273,9 +1295,8 @@ class MLFlowRepository(SientiaMonitoring):
             end_time = datetime.now()
 
             if isinstance(predict_data, pd.DataFrame):
-                self.debug(
+                self._debug_dataframe(
-                    f'Data received from model prediction: {predict_data.to_dict(orient="records")}',
+                    'Data received from model prediction:', predict_data, metadata
-                    metadata,
                 )
 
                 # predict_data.to_csv(
@@ -1283,6 +1304,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)
@@ -1290,7 +1315,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 {
@@ -1344,7 +1369,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)
 

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,

laborious/worker/worker.py

@@ -152,8 +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,
@@ -199,11 +198,11 @@ async def main():
                 activities.format_transformed_data,
                 activities.format_prediction,
                 activities.format_default_prediction,
-                activities.get_last_timestamp,
                 # OPC
                 activities.write_opc_data,
-                # Postgres
+                # Postgres / MinIO offload
-                activities.load_custom_query,
+                activities.load_query_with_minio_offload,
+                activities.cleanup_minio_objects_expired,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres,
                 activities.write_metrics,
@@ -220,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():

laborious/workflows/drift.py

@@ -44,9 +44,9 @@ 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(
+        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,

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')
@@ -73,26 +74,26 @@ 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']:
+        storage_payload = MinioDataFramePayload.from_dict(storage_result)
-            return
+        if not storage_payload.has_data():
+            raise ValueError('No data returned from query')
 
         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,
             },

laborious/workflows/predictions_batch.py

@@ -83,13 +83,14 @@ class PredictionsBatch:
             }
         }
 
-        # Load data using custom query
+        # Load data using custom query with optional MinIO offload for large frames
-        data = await workflow.execute_local_activity_method(
+        data = await workflow.execute_activity_method(
-            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'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=300),
@@ -104,16 +105,19 @@ 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']),
             '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),

laborious/workflows/simple_metrics.py

@@ -45,9 +45,9 @@ 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(
+        target_data = await workflow.execute_activity_method(
             Activities.load_custom_query,
             {
                 **metadata,

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'],
@@ -196,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', 'error'),
+                'unique_columns': ['model_id', 'timestamp'],
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=180),

laborious/workflows/sub_workflows/prediction_process.py

@@ -87,13 +87,42 @@ class PredictionProcess:
         model_config = input_data.get('model_config', {})
         save_transform = input_data.get('save_transform', True)
 
-        # Get last timestamp for incremental processing
+        try:
-        last_timestamp = await workflow.execute_local_activity_method(
+            await self._run_prediction_pipeline(
-            Activities.get_last_timestamp,
+                input_data,
-            {**metadata, 'data': data},
+                metadata,
-            retry_policy=retry_policy,
+                data,
-            start_to_close_timeout=timedelta(minutes=1),
+                model_id,
-        )
+                model_name,
+                model_config,
+                save_transform,
+            )
+            await workflow.execute_activity_method(
+                Activities.cleanup_minio_objects_expired,
+                {**metadata, 'data': data},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+        except Exception as e:
+            await workflow.execute_activity_method(
+                Activities.cleanup_minio_objects_expired,
+                {**metadata, 'data': data},
+                retry_policy=retry_policy,
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+            raise e
+
+    async def _run_prediction_pipeline(
+        self,
+        input_data: dict[str, Any],
+        metadata: dict[str, Any],
+        data: dict[str, Any],
+        model_id: str,
+        model_name: str,
+        model_config: dict[str, Any],
+        save_transform: bool,
+    ) -> None:
+        last_timestamp = data['last_timestamp']
 
         # Apply input data quality gates
         gate_input = {
@@ -117,7 +146,7 @@ class PredictionProcess:
             return
 
         # Request MLFlow model transformation
-        response_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,
@@ -130,7 +159,7 @@ class PredictionProcess:
             {
                 **metadata,
                 'filters': input_data['mlflow_transform_filters'],
-                'data': response_data,
+                'data': transformed_data,
                 'type': 'transform',
                 'path_priority': input_data['path_priority'],
             },
@@ -144,8 +173,6 @@ class PredictionProcess:
         ):
             return
 
-        transformed_data = response_data['content']
-
         path_flag, confidence, comment = await workflow.execute_local_activity_method(
             Activities.mlflow_content_gate,
             {
@@ -164,7 +191,7 @@ class PredictionProcess:
         ):
             return
 
-        response_data = await workflow.execute_local_activity_method(
+        predicted_data = await workflow.execute_activity_method(
             Activities.request_predict,
             {
                 **metadata,
@@ -182,7 +209,7 @@ class PredictionProcess:
             {
                 **metadata,
                 'filters': input_data['mlflow_predict_filters'],
-                'data': response_data,
+                'data': predicted_data,
                 'type': 'predict',
                 'path_priority': input_data['path_priority'],
             },
@@ -201,8 +228,9 @@ class PredictionProcess:
             'subworkflow.format_and_export_prediction',
             {
                 'metadata': metadata,
+                'on_conflict': input_data.get('on_conflict', 'error'),
                 '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,
@@ -221,7 +249,7 @@ class PredictionProcess:
 
     async def path_flag_handler(
         self,
-        data: dict,
+        data: dict[str, Any],
         path_flag: str,
         input_data: dict,
         confidence: int,
@@ -310,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

pyproject.toml

@@ -23,6 +23,7 @@ exclude = [
     "*.pyc",
     ".pytest_cache",
     "htmlcov",
+    "tests/laborious/workflows/subworkflows/test_prediction_process.py",
 ]
 
 [tool.ruff.lint]

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

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

requirements.txt

@@ -3,8 +3,8 @@ 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.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
@@ -13,4 +13,6 @@ pyarrow
 kaleido
 hyperopt
 shap
-pycurl
+pycurl
+scipy<1.14.0
+scikit-learn==1.5.2

tests/conftest.py

@@ -1,3 +1,49 @@
+import os
+import sys
+from unittest.mock import MagicMock
+
+# 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_MEGABYTES', '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

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,
@@ -41,8 +43,9 @@ def test___init__(
         'endpoint_url': 'localhost:9000',
         'access_key': 'minio',
         'secret_key': 'minio123',
-        'region_name': 'us-east-1',
         'default_bucket': 'test',
+        'retention_hours': 24,
+        'secure': False,
     }
 
     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,17 @@ def test___init__(
         metrics_controller=mock_metrics_controller.return_value,
     )
 
+    mock_minio_repository.assert_called_once_with(
+        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'],
+    )
+
 
 @mark.asyncio
 @patch('laborious.activities.activities.Storage')
@@ -147,7 +163,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,
@@ -170,8 +188,9 @@ 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,
+        'secure': False,
     }
 
     mlflow_config = {'host': 'localhost', 'port': 5000, 'username': 'mlflow', 'password': 'mlflow'}

tests/laborious/activities/test_api.py

@@ -61,6 +61,70 @@ 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',
+            'runtime': 'local',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write_pi_web_api_data',
+        },
+    ):
+        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',
+    }
+
+
+@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',
+            'runtime': 'k8s',
+            'model_name': 'test_model',
+            'workflow_name': 'test_workflow',
+            'operation_type': 'write',
+        },
+    ):
+        labels = api_instance.get_pi_web_api_core_labels(
+            metadata=metadata['metadata'], operation_type='write'
+        )
+    assert labels['operation_type'] == 'write'
+    assert labels['runtime'] == 'k8s'
+
+
 def test__init__():
     api = API(
         base_url='https://test-pi-server.com',
@@ -98,6 +162,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',
         }
@@ -280,6 +345,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',
     }
@@ -313,6 +379,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',
     }
@@ -341,6 +408,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',
     }
@@ -373,6 +441,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',
     }
@@ -401,6 +470,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',
     }

tests/laborious/activities/test_gates.py

@@ -1,11 +1,37 @@
 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
 
 
+@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.
+
+    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 +66,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'],
     }
 
@@ -64,8 +90,8 @@ 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': {'value': []},
+        'data': _minio_payload(DataFrame({'value': []})),
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
@@ -77,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,
@@ -90,7 +116,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'],
     }
 
@@ -107,8 +133,8 @@ 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': {'value': []},
+        'data': _minio_payload(DataFrame({'value': []})),
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
 
@@ -121,12 +147,46 @@ 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': {'value': [1, 2, 3]},
+        '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'],
     }
 
@@ -144,7 +204,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 +232,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 +261,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'],
     }
@@ -213,11 +282,11 @@ 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': {
+        'data': _minio_payload(
-            'success': False,
+            {'content': {'message': 'API error occurred', 'traceback': 'error trace'}},
-            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
+            status={'success': False, 'message': 'API error occurred'},
-        },
+        ),
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -231,16 +300,37 @@ 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': {
+        'data': _minio_payload(
-            'success': True,
+            {'content': {'message': 'success'}},
-            'content': {'message': 'success'},
+            status={'success': True},
-        },
+        ),
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -259,10 +349,7 @@ async def test_mlflow_content_gate_invalid_filter(gates_activity):
     input_data = {
         **metadata,
         'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
-        'data': {
+        'data': _minio_payload(DataFrame({'value': [1, 2, 3]})),
-            'success': True,
-            'content': {'message': 'success'},
-        },
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -287,10 +374,7 @@ async def test_mlflow_content_gate_filter_exception(
     input_data = {
         **metadata,
         'filters': {'API_ERROR': {'POLICY': 'STOP'}},
-        'data': {
+        'data': _minio_payload(DataFrame({'value': [1, 2, 3]})),
-            'success': False,
-            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
-        },
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -317,7 +401,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'],
     }
@@ -335,8 +419,8 @@ 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': {'value': [None, None, None]},
+        'data': _minio_payload(DataFrame({'value': [None, None, None]})),
         'type': 'test',
         'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
     }
@@ -356,7 +440,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 +453,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,13 +530,18 @@ async def test_format_prediction_no_timestamp(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': _minio_payload(
-            'prediction': {'2023-05-26 11:12:27': 1},
+            DataFrame(
-            'response_time': {'2023-05-26 11:12:27': 0.1},
+                {
-        },
+                    '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',
+        'timestamp': '2023-05-26 11:12:27',
     }
 
     # Act
@@ -457,21 +562,26 @@ async def test_format_prediction_with_timestamp_erl(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': _minio_payload(
-            'prediction': {
+            DataFrame(
-                '2023-05-26 11:12:27': 1,
+                {
-                '2023-05-26 11:12:28': 2,
+                    'prediction': {
-                '2023-05-26 11:12:29': 3,
+                        '2023-05-26 11:12:27': 1,
-            },
+                        '2023-05-26 11:12:28': 2,
-            'response_time': {
+                        '2023-05-26 11:12:29': 3,
-                '2023-05-26 11:12:27': 0.1,
+                    },
-                '2023-05-26 11:12:28': 0.2,
+                    'response_time': {
-                '2023-05-26 11:12:29': 0.3,
+                        '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',
+        'timestamp': '2023-05-26 11:12:27',
     }
 
     # Act
@@ -492,21 +602,26 @@ async def test_format_prediction_with_timestamp_lts(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': _minio_payload(
-            'prediction': {
+            DataFrame(
-                '2023-05-26 11:12:27': 1,
+                {
-                '2023-05-26 11:12:28': 2,
+                    'prediction': {
-                '2023-05-26 11:12:29': 3,
+                        '2023-05-26 11:12:27': 1,
-            },
+                        '2023-05-26 11:12:28': 2,
-            'response_time': {
+                        '2023-05-26 11:12:29': 3,
-                '2023-05-26 11:12:27': 0.1,
+                    },
-                '2023-05-26 11:12:28': 0.2,
+                    'response_time': {
-                '2023-05-26 11:12:29': 0.3,
+                        '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',
+        'timestamp': '2023-05-26 11:12:27',
     }
 
     # Act
@@ -527,14 +642,23 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
     # Arrange
     input_data = {
         **metadata,
-        'data': {
+        'data': _minio_payload(
-            'prediction': [1, 2, 3],
+            DataFrame(
-            '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'],
+                    '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',
+        'timestamp': '2023-05-26 11:12:27',
     }
     gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))
 
@@ -547,76 +671,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': {
+        'data': _minio_payload(
-            'var1': {'2023-05-26 11:12:27': 1.0},
+            DataFrame(
-            'var2': {'2023-05-26 11:12:27': 2.0},
+                {
-        },
+                    '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 is payload_result
-    assert result['variable'] == {0: 'var1', 1: 'var2'}
+    mock_from_dataframe.assert_called_once()
-    assert result['value'] == {0: 1.0, 1: 2.0}
+    kwargs = mock_from_dataframe.call_args.kwargs
-    assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
+    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': {
+        'data': _minio_payload(
-            'var1': {
+            DataFrame(
-                '2023-05-26 11:12:27': 1.0,
+                {
-                '2023-05-26 11:12:28': 2.0,
+                    'var1': {
-            },
+                        '2023-05-26 11:12:27': 1.0,
-            'var2': {
+                        '2023-05-26 11:12:28': 2.0,
-                '2023-05-26 11:12:27': 3.0,
+                    },
-                '2023-05-26 11:12:28': 4.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 result is payload_result
-    assert len(result['variable']) == 4
+    mock_from_dataframe.assert_called_once()
-    assert len(result['value']) == 4
+    kwargs = mock_from_dataframe.call_args.kwargs
-    assert len(result['model_id']) == 4
+    assert kwargs['model_name'] == 'test_model'
-    assert all(v == 'test_model' for v in result['model_id'].values())
+    assert kwargs['operation'] == 'transform'
-    assert set(result['variable'].values()) == {'var1', 'var2'}
+    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 is payload_result
-    assert result['variable'] == {}
+    mock_from_dataframe.assert_called_once()
-    assert result['value'] == {}
+    kwargs = mock_from_dataframe.call_args.kwargs
-    assert result['model_id'] == {}
+    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 +865,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):
@@ -750,15 +879,18 @@ 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,
+        'operation_type': 'predict',
+        '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={
+                tags=core_tags,
-                    'pod_id': gates_activity.pod_id,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
-                },
             ),
         ]
     )
@@ -767,11 +899,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
             call(
                 metric_object=mock_metrics.PREDICTION_CONFIDENCE_MONITOR,
                 method='set',
-                tags={
+                tags=core_tags,
-                    'pod_id': gates_activity.pod_id,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
-                },
                 value=0.9,
             ),
         ]
@@ -781,11 +909,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
             call(
                 metric_object=mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR,
                 method='observe',
-                tags={
+                tags=core_tags,
-                    'pod_id': gates_activity.pod_id,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
-                },
                 value=0.1,
             ),
         ]
@@ -795,9 +919,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,
+                    **core_tags,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
                     'opc_server_id': 'server1',
                     'tag': 'tag1',
                 },
@@ -810,9 +932,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,
+                    **core_tags,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
                     'opc_server_id': 'server1',
                     'tag': 'tag1',
                 },
@@ -825,9 +945,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,
+                    **core_tags,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
                     'opc_server_id': 'server1',
                     'tag': 'tag2',
                 },
@@ -840,9 +958,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,
+                    **core_tags,
-                    'model_name': metadata['metadata']['model_name'],
-                    'workflow_name': metadata['metadata']['workflow_name'],
                     'opc_server_id': 'server1',
                     'tag': 'tag2',
                 },
@@ -873,6 +989,8 @@ 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,
+            'operation_type': 'predict',
             'model_name': metadata['metadata']['model_name'],
             'workflow_name': metadata['metadata']['workflow_name'],
             'opc_server_id': 'server1',

tests/laborious/activities/test_mlflow.py

@@ -8,24 +8,38 @@ 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):
+    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={
+        minio_repository=minio_repo,
-            'endpoint_url': 'http://localhost:9000',
+        logger=logger,
-            'access_key': 'minio',
+        notification_handler=notification_handler,
-            'secret_key': 'minio123',
+        metrics_controller=metrics_controller,
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-        metrics_controller=AsyncMock(),
     )
 
     assert mlflow.mlflow_host == 'http://localhost'
@@ -38,14 +52,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',
     )
 
 
@@ -53,21 +66,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={
+        minio_repository=minio_repo,
-            'endpoint_url': 'http://localhost:9000',
+        logger=logger,
-            'access_key': 'minio',
+        notification_handler=notification_handler,
-            'secret_key': 'minio123',
+        metrics_controller=metrics_controller,
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
-        logger=MagicMock(),
-        notification_handler=MagicMock(),
-        metrics_controller=AsyncMock(),
     )
 
     mlflow.model_monitoring_repository = AsyncMock()
@@ -96,139 +115,147 @@ metadata = {
 
 
 @mark.asyncio
-@patch('laborious.activities.mlflow.DataFrame')
+@patch(
-@patch('laborious.activities.mlflow.max')
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
-async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
+    new_callable=AsyncMock,
-    mock_max.return_value = '2024-01-02'
+)
-    # Mock input data
+async def test_request_transform_success(mock_from_dataframe, mlflow):
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
     input_data = {
         **metadata,
-        'data': [
+        'data': payload,
-            {
-                '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',
-            },
-        ],
         'model_name': 'test_model',
         'model_config': {},
     }
 
-    # Mock the transform response
+    transform_response = {'success': True, 'content': MagicMock()}
-    expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
-    mlflow.model_monitoring_repository.transform.return_value = expected_response
 
-    mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
+    data_mock.sort_values.return_value = data_mock
-    mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value
+    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(
-        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)
-    # mock_dataframe.reset_index.assert_called_once()
-    mock_dataframe.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']
     )
+    mock_from_dataframe.assert_called_once()
+    assert response_data == mock_from_dataframe.return_value
 
 
 @mark.asyncio
-@patch('laborious.activities.mlflow.DataFrame')
+@patch(
-@patch('laborious.activities.mlflow.to_datetime')
+    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
-@patch('laborious.activities.mlflow.max')
+    new_callable=AsyncMock,
-async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
+)
-    mock_max.return_value = '2024-01-02'
+async def test_request_transform_failure(mock_from_dataframe, mlflow):
-    # Mock input data
+    data_mock = MagicMock()
+    payload = AsyncMock()
+    payload.retrieve = AsyncMock(return_value=data_mock)
+
     input_data = {
         **metadata,
-        'data': {
+        'data': payload,
-            '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},
-        },
         'model_name': 'test_model',
         'model_config': {},
     }
 
-    # Mock the predict response
+    transform_response = {'success': False, 'message': 'Transform failed'}
-    expected_response = {'prediction': [0.5, 0.6]}
+    mlflow.model_monitoring_repository.transform.return_value = transform_response
-    mlflow.model_monitoring_repository.predict.return_value = expected_response
+
+    data_mock.sort_values.return_value = data_mock
+    data_mock.drop_duplicates.return_value = data_mock
+    data_mock.pivot.return_value = data_mock
+
+    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'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    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(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
 
-    # Call the method
     response_data = await mlflow.request_predict(input_data)
 
-    mock_dataframe.assert_called_once_with(input_data['data'])
+    data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
-    mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
-    mock_dataframe.return_value.__setitem__.assert_any_call(
-        'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
-    )
-    mock_dataframe.return_value.__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
-    )
-    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', mock_dataframe.return_value, {}, metadata['metadata']
+        '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'],
+        last_timestamp=payload.last_timestamp,
+        logger=mlflow.logger,
+    )
+    assert response_data == mock_from_dataframe.return_value
 
 
 @mark.asyncio
@@ -240,12 +267,14 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
         'message': 'Model retrained successfully.',
     }
 
-    mlflow.minio_repository.get_parquet_as_dataframe.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',
@@ -255,8 +284,6 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
         }
     )
 
-    raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
-
     timestamp = raw_data.__getitem__.return_value.max.return_value
 
     raw_data.sort_values.assert_not_called()
@@ -309,21 +336,21 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
 
 @mark.asyncio
 @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, mlflow):
     mlflow.model_monitoring_repository.retrain_model.return_value = {
-        'success': False,
+        'success': True,
-        'traceback': 'test_traceback',
+        'experiment': 'test_experiment',
-        'message': 'Model retrained failed.',
+        'message': 'Model retrained successfully.',
     }
 
-    mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock(
+    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
-        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',
@@ -333,7 +360,35 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow)
         }
     )
 
-    raw_data = mlflow.minio_repository.get_parquet_as_dataframe.return_value
+    assert response['success'] is True
+    mlflow.minio_repository.download_file.assert_not_called()
+
+
+@mark.asyncio
+@patch('laborious.activities.mlflow.to_datetime')
+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.',
+    }
+
+    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': payload,
+            'model_name': 'test_model',
+            'model_config': {
+                'target': 'target',
+                'transform_flavor': 'sklearn',
+                'predict_flavor': 'pyfunc',
+            },
+        }
+    )
 
     timestamp = raw_data.__getitem__.return_value.max.return_value
 
@@ -398,14 +453,9 @@ 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(
-        'Error loading retrain data'
-    )
-
     response = await mlflow.retrain_model(
         {
             **metadata,
-            'object_key': 'test_object_key',
             'model_name': 'test_model',
             'model_config': {
                 'target': 'target',
@@ -417,7 +467,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,
     }

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'

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
@@ -7,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',
@@ -28,13 +38,8 @@ def storage(mock_minio_repository):
         dbname='postgres',
         min_connections=1,
         max_connections=10,
-        minio_config={
+        retention_hours=24,
-            'endpoint_url': 'localhost:9000',
+        minio_repository=mock_minio_repository.return_value,
-            'access_key': 'minio',
-            'secret_key': 'minio123',
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
         logger=MagicMock(),
         notification_handler=MagicMock(),
         metrics_controller=AsyncMock(),
@@ -46,6 +51,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,
@@ -54,29 +60,16 @@ def test___init___not_hasattr(mock_minio_repository):
         dbname='postgres',
         min_connections=1,
         max_connections=10,
-        minio_config={
+        retention_hours=24,
-            'endpoint_url': 'localhost:9000',
+        minio_repository=minio_repo,
-            'access_key': 'minio',
-            'secret_key': 'minio123',
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
         logger=logger,
         notification_handler=notification_handler,
         metrics_controller=metrics_controller,
     )
     assert isinstance(storage, Postgres)
 
-    mock_minio_repository.assert_called_once_with(
+    assert storage.minio_repository is minio_repo
-        logger=logger,
+    mock_minio_repository.assert_not_called()
-        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,
-    )
 
 
 @patch('laborious.activities.storage.MinioRepository')
@@ -93,28 +86,15 @@ def test___init___none_minio_repository(mock_minio_repository, storage):
         dbname='postgres',
         min_connections=1,
         max_connections=10,
-        minio_config={
+        retention_hours=24,
-            'endpoint_url': 'localhost:9000',
+        minio_repository=None,
-            'access_key': 'minio',
-            'secret_key': 'minio123',
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
         logger=logger,
         notification_handler=notification_handler,
         metrics_controller=metrics_controller,
     )
 
-    mock_minio_repository.assert_called_once_with(
+    assert storage.minio_repository is None
-        logger=logger,
+    mock_minio_repository.assert_not_called()
-        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,
-    )
 
 
 @patch('laborious.activities.storage.MinioRepository')
@@ -127,13 +107,8 @@ def test___init___done_repository(mock_minio_repository, storage):
         dbname='postgres',
         min_connections=1,
         max_connections=10,
-        minio_config={
+        retention_hours=24,
-            'endpoint_url': 'localhost:9000',
+        minio_repository=mock_minio_repository.return_value,
-            'access_key': 'minio',
-            'secret_key': 'minio123',
-            'region_name': 'us-east-1',
-            'default_bucket': 'test',
-        },
         logger=MagicMock(),
         notification_handler=MagicMock(),
         metrics_controller=AsyncMock(),
@@ -142,72 +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.store_dataframe_as_parquet = AsyncMock()
-    storage.minio_repository.minio_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',
-        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'
-
-
-@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'})
-    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()
 
@@ -222,3 +131,193 @@ def test___del__(storage):
     storage.__del__()
 
     storage.close.assert_called_once()
+
+
+@mark.asyncio
+async def test_load_query_with_minio_offload_no_rows(storage):
+    storage.load_custom_query = AsyncMock(return_value=None)
+    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}])
+    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
+async def test_load_query_with_minio_offload_minio(storage):
+    storage.load_custom_query = AsyncMock(return_value=[{'a': 1}])
+    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 == storage_result
+    mock_from_dataframe.assert_awaited_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()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    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
+
+    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, 'data': data_mock})
+
+
+@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()
+
+    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
+    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()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    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()
+
+    data_mock = MagicMock()
+    data_mock.cleanup_prefix.return_value = 'training_datasets/m'
+    result = await storage.cleanup_minio_objects_expired({**metadata, 'data': data_mock})
+
+    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()

tests/laborious/utils/models/test_minio_dataframe_payload.py

@@ -0,0 +1,266 @@
+from datetime import datetime
+from io import BytesIO
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from pandas import DataFrame
+
+from laborious.utils.models.minio_dataframe_payload import (
+    MinioDataFramePayload,
+    _build_object_key,
+)
+
+
+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_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_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_when_inline_returns_none():
+    payload = MinioDataFramePayload(last_timestamp='t', data={'x': [1]}, object_key=None)
+    assert MinioDataFramePayload.cleanup_prefix(payload) is None
+
+
+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
+
+
+@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 == '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 == 'prediction_datasets/my-model'
+    assert key.startswith('prediction_datasets/my-model/')
+
+
+def test_estimate_size_bytes_fallback():
+    df = DataFrame({'a': [1, 2]})
+    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
+
+
+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()
+
+
+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

tests/laborious/utils/repository/test_minio_repository.py

@@ -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()

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 == {
+    assert output['success'] is True
-        'success': True,
+    assert output['content'] is mlflow_repository.detect_and_parse_datetime_index.return_value
-        'content': mlflow_repository.detect_and_parse_datetime_index.return_value.to_dict.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'] == {
+    content = output['content']
-        'prediction': {'index_1': 2, 'index_2': 3},
+    assert isinstance(content, DataFrame)
-        'response_time': {'index_1': ANY, 'index_2': ANY},
+    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'] == {
+    content = output['content']
-        'prediction': {'index_1': 2, 'index_2': 3},
+    assert isinstance(content, DataFrame)
-        'response_time': {'index_1': ANY, 'index_2': ANY},
+    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()

tests/laborious/utils/test_connectors_config.py

@@ -100,8 +100,9 @@ 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,
+        'secure': False,
     }
 
 
@@ -115,6 +116,7 @@ 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,
+        'secure': False,
     }

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,
@@ -91,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,
@@ -98,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,
@@ -141,6 +145,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 +181,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 +195,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 +208,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 +218,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,
@@ -238,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,
@@ -245,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,
@@ -287,6 +296,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 +321,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,
@@ -343,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,
@@ -389,6 +401,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 +428,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,
@@ -448,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,
@@ -455,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,
@@ -496,6 +512,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',
@@ -551,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,
@@ -558,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,
@@ -597,6 +616,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',
@@ -638,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,
@@ -645,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,

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
 
@@ -17,6 +17,7 @@ metadata = {
         'model_name': 'test_model',
         'workflow_name': 'test_workflow',
         'schema_name': 'test_schedule',
+        'schedule_name': 'test_schedule',
     },
 }
 
@@ -26,9 +27,14 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
     input_data = {
         'metadata': metadata,
-        'data': {'test': 'data'},
+        'data': data_payload,
         'schema': 'test_schema',
         'table_name': 'test_table',
         'transform_table_name': 'test_transform_table',
@@ -45,15 +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 = [
-        '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)
         ('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'),
     ]
@@ -62,21 +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 == 7
+    assert workflow_mock.execute_local_activity_method.call_count == 4
-
-    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,
-            )
-        ]
-    )
     workflow_mock.execute_local_activity_method.assert_has_calls(
         [
             call(
@@ -92,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 +124,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'],
                 },
@@ -139,13 +133,13 @@ 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,
                 {
                     **metadata,
-                    'data': 'transformed_data',
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
                 },
@@ -175,9 +169,10 @@ async def test_run(workflow_mock, prediction_process):
         'subworkflow.format_and_export_prediction',
         {
             'metadata': metadata,
+            'on_conflict': 'error',
             'path_flag': 'continue',
-            'data': 'predicted_data',
+            'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'},
-            'transformed_data': 'transformed_data',
+            'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
             'prediction_confidence': 0.95,
             'timestamp': '2024-01-01',
             'model_id': 1,
@@ -199,9 +194,14 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
     input_data = {
         'metadata': metadata,
-        'data': {'test': 'data'},
+        'data': data_payload,
         'schema': 'test_schema',
         'table_name': 'test_table',
         'transform_table_name': 'test_transform_table',
@@ -217,7 +217,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
     ]
 
@@ -225,18 +224,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,
                 {
@@ -258,9 +248,14 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
     input_data = {
         'metadata': metadata,
-        'data': {'test': 'data'},
+        'data': data_payload,
         'schema': 'test_schema',
         'table_name': 'test_table',
         'transform_table_name': 'test_transform_table',
@@ -275,10 +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 = [
-        '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)
     ]
 
@@ -286,20 +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 == 4
+    assert workflow_mock.execute_local_activity_method.call_count == 2
-    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(
@@ -315,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,
@@ -354,9 +338,14 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
     input_data = {
         'metadata': metadata,
-        'data': {'test': 'data'},
+        'data': data_payload,
         'schema': 'test_schema',
         'table_name': 'test_table',
         'transform_table_name': 'test_transform_table',
@@ -371,10 +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 = [
-        '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)
         ('continue', 0.95, 'Error'),
         # mlflow_content_gate (transform)
@@ -385,21 +376,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 == 3
 
-    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(
@@ -415,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,
@@ -452,7 +430,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,
@@ -470,9 +448,14 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
     input_data = {
         'metadata': metadata,
-        'data': {'test': 'data'},
+        'data': data_payload,
         'schema': 'test_schema',
         'table_name': 'test_table',
         'transform_table_name': 'test_transform_table',
@@ -487,15 +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 = [
-        '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)
         ('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)
     ]
 
@@ -503,20 +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 == 7
+    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(
@@ -532,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,
@@ -569,7 +541,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,
@@ -579,12 +551,12 @@ 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,
                 {
-                    'data': 'transformed_data',
+                    'data': {'content': 'transformed_data', 'timestamp': '2024-01-01'},
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
                     **metadata,
@@ -768,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',
         },
     )
 
@@ -816,3 +789,54 @@ 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.__getitem__ = (
+        lambda self, key: '2024-01-01' if key == 'last_timestamp' else MagicMock()
+    )
+    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_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'),
+        ('continue', 0.95, ''),
+        ('continue', 0.95, ''),
+        ('continue', 0.95, ''),
+    ]
+
+    await prediction_process.run(input_data)
+
+    workflow_mock.execute_activity_method.assert_any_call(
+        Activities.cleanup_minio_objects_expired,
+        {**metadata, 'data': data_payload},
+        retry_policy=ANY,
+        start_to_close_timeout=ANY,
+    )

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.start_activity_method.side_effect = [target_data, reference_data]
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
-    workflow_mock.execute_local_activity_method.return_value = drift_data
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # 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.start_activity_method.side_effect = [target_data, reference_data]
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
-    workflow_mock.execute_local_activity_method.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.start_activity_method.side_effect = [target_data, reference_data]
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=drift_data)
-    workflow_mock.execute_local_activity_method.return_value = drift_data
+    workflow_mock.execute_activity_method = AsyncMock(return_value=None)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # Act
     await drift.run(input_data)

tests/laborious/workflows/test_minimal_retrain.py

@@ -39,9 +39,19 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
         },
     }
 
+    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=[
-            {'success': True, 'object_key': 'test_object_key'},
+            storage_result,
             {'success': True, 'experiment': 'test_experiment'},
             {
                 'success': True,
@@ -58,13 +68,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 +87,7 @@ async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
                 Activities.retrain_model,
                 {
                     **metadata,
-                    'object_key': 'test_object_key',
+                    'data': storage_result,
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
                 },
@@ -161,9 +170,19 @@ async def test_run_storage_fail(workflow_mock: AsyncMock, minimal_retrain: Minim
         },
     }
 
+    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=[
-            {'success': False, 'object_key': 'test_object_key'},
+            storage_result,
             {'success': True, 'experiment': 'test_experiment'},
             {
                 'success': True,
@@ -175,16 +194,18 @@ 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.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,
@@ -211,9 +232,19 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
         },
     }
 
+    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=[
-            {'success': True, 'object_key': 'test_object_key'},
+            storage_result,
             {'success': False, 'experiment': 'test_experiment'},
             {
                 'success': True,
@@ -230,13 +261,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 +280,7 @@ async def test_run_fail_retrain(workflow_mock: AsyncMock, minimal_retrain: Minim
                 Activities.retrain_model,
                 {
                     **metadata,
-                    'object_key': 'test_object_key',
+                    'data': storage_result,
                     'model_name': input_data['model_name'],
                     'model_config': input_data['model_config'],
                 },

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
 
@@ -12,19 +12,19 @@ def predictions_batch() -> PredictionsBatch:
 
 
 metadata = {
-    'metadata': {
+    'model_id': 'test_model_id',
-        'model_id': 'test_model_id',
+    'model_name': 'test_model',
-        'model_name': 'test_model',
+    'workflow_name': 'predictions_batch',
-        'workflow_name': 'predictions_batch',
+    'schedule_name': 'test_schedule',
-        'schedule_name': 'test_schedule',
-    },
 }
 
 
 @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'}
+    activity_return = MagicMock()
+    workflow_mock.execute_activity_method.return_value = activity_return
+
     input_data = {
         'schedule_name': 'test_schedule',
         'model_name': 'test_model',
@@ -42,14 +42,15 @@ 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,
+                    'metadata': metadata,
                     'query': input_data['query'],
                     'datetime_columns': input_data.get('datetime_columns', []),
+                    'model_name': input_data['model_name'],
                 },
                 retry_policy=ANY,
                 start_to_close_timeout=ANY,
@@ -57,23 +58,44 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch
         ]
     )
     prediction_input = {
-        'metadata': metadata,
+        'metadata': {'metadata': metadata},
-        'data': {'data': 'test_data'},
+        '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']),
         '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),

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(side_effect=[target_data, None])
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # 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'],
+                    'data': simple_metrics_data,
-                    'target_data': target_data,
+                    'schema': input_data['schema'],
-                    'metrics': input_data['metrics'],
+                    'table_name': input_data['target_table_name'],
-                    'interval_minutes': input_data['interval_minutes'],
+                    'timestamp_conversion': {
+                        'column': 'timestamp',
+                        'format': DATETIME_FORMAT_WITH_TZ,
+                    },
                 },
                 retry_policy=ANY,
                 start_to_close_timeout=ANY,
             ),
         ]
     )
-
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
-    # Assert - Check export_data_to_postgres call
+        Activities.calculate_simple_metrics,
-    workflow_mock.execute_activity_method.assert_called_once_with(
-        Activities.export_data_to_postgres,
         {
             **metadata,
-            'data': simple_metrics_data,
+            'model_id': input_data['model_id'],
-            'schema': input_data['schema'],
+            'target_data': target_data,
-            'table_name': input_data['target_table_name'],
+            'metrics': input_data['metrics'],
-            'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
+            '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(return_value=target_data)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # 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
+    assert workflow_mock.execute_activity_method.call_count == 1
-    workflow_mock.execute_activity_method.assert_not_called()
 
 
 @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(return_value=target_data)
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # 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_called_once()
-    workflow_mock.execute_activity_method.assert_not_called()
+    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(side_effect=[target_data, None])
-
+    workflow_mock.execute_local_activity_method = AsyncMock(return_value=simple_metrics_data)
-    workflow_mock.execute_activity_method = AsyncMock()
 
     # 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(
+    workflow_mock.execute_activity_method.assert_any_call(
-        [
+        Activities.load_custom_query,
-            call(
+        ANY,
-                Activities.load_custom_query,
+        retry_policy=ANY,
-                ANY,
+        start_to_close_timeout=ANY,
-                retry_policy=ANY,
+    )
-                start_to_close_timeout=ANY,
+    workflow_mock.execute_local_activity_method.assert_called_once_with(
-            ),
+        Activities.calculate_simple_metrics,
-            call(
+        {
-                Activities.calculate_simple_metrics,
+            **metadata,
-                {
+            'model_id': input_data['model_id'],
-                    **metadata,
+            'target_data': target_data,
-                    'model_id': input_data['model_id'],
+            'metrics': ['rmse', 'mse', 'mae', 'r2'],  # Default value
-                    'target_data': target_data,
+            'interval_minutes': input_data['interval_minutes'],
-                    'metrics': ['rmse', 'mse', 'mae', 'r2'],  # Default value
+        },
-                    'interval_minutes': input_data['interval_minutes'],
+        retry_policy=ANY,
-                },
+        start_to_close_timeout=ANY,
-                retry_policy=ANY,
-                start_to_close_timeout=ANY,
-            ),
-        ]
     )

validate.sh

@@ -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

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"
 
@@ -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
@@ -222,15 +222,69 @@ 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
-    value: "FvcxOPX55j"
+    value: "LiArt4eNmJ"
-  - name: MINIO_REGION_NAME
-    value: "sa-east-1"
   - name: MINIO_DEFAULT_BUCKET
     value: "sientia"
+  - name: MINIO_RETENTION_HOURS
+    value: "24"
+  - name: SIENTIA_MINIO_OFFLOAD_THRESHOLD_MEGABYTES
+    value: "0.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: "3"
+  # Small initial poller count warms up gradually instead of spiking task fetches.
+  - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_INITIAL
+    value: "5"
+  # Cap workflow pollers to limit scheduling pressure and avoid over-polling.
+  - name: PREDICTIONSBATCH_WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM
+    value: "15"
+  # Keep at least two activity pollers so activity queues do not starve during spikes.
+  - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MINIMUM
+    value: "3"
+  # Moderate initial activity pollers for faster ramp-up with controlled pressure.
+  - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_INITIAL
+    value: "10"
+  # Limit max activity pollers to preserve CPU for workflow-task completion.
+  - name: PREDICTIONSBATCH_ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM
+    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"