vitor-aignosi
892823df11
Enhance Activities and Prediction Workflows with PI Web API Integration - Updated the Activities class to include API integration, allowing for configuration of PI Web API parameters. - Modified prediction workflows to support output configuration for PI Web API, enabling data writing to the API. - Refactored connectors_config.py by removing unused PostgreSQL and MongoDB configuration functions. - Added tests to validate the new PI Web API functionality in activities and workflows, ensuring robust integration and functionality.
121 lines
5.3 KiB
Python
121 lines
5.3 KiB
Python
from temporalio import workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from datetime import timedelta
|
|
from typing import Any
|
|
|
|
from sientia_do.temporal.policies import retry_policy
|
|
|
|
from laborious.activities.activities import Activities
|
|
|
|
|
|
@workflow.defn(name='predictions_batch')
|
|
class PredictionsBatch:
|
|
"""
|
|
Main batch prediction workflow for the Laborious system.
|
|
|
|
This workflow orchestrates the complete batch prediction process, handling
|
|
data loading, configuration management, and workflow delegation. It serves
|
|
as the primary entry point for batch prediction operations and ensures
|
|
proper data preparation before ML model inference.
|
|
|
|
The workflow implements a robust data processing pipeline with:
|
|
- Custom SQL query execution for data loading
|
|
- Comprehensive configuration management
|
|
- Data quality filter application
|
|
- MLFlow model integration
|
|
- Workflow delegation to specialized sub-workflows
|
|
|
|
Workflow Execution:
|
|
1. Data Loading: Executes custom SQL query to load prediction data
|
|
2. Configuration Preparation: Sets up prediction parameters and filters
|
|
3. Workflow Delegation: Spawns PredictionProcess child workflow
|
|
4. Error Handling: Implements comprehensive error handling and retry policies
|
|
"""
|
|
|
|
@workflow.run
|
|
async def run(self, input_data: dict[str, Any]):
|
|
"""
|
|
Execute the batch prediction workflow.
|
|
|
|
This method orchestrates the complete batch prediction process by:
|
|
1. Loading data using the provided custom SQL query
|
|
2. Preparing prediction configuration and filters
|
|
3. Delegating to the PredictionProcess workflow for ML operations
|
|
|
|
The method implements comprehensive error handling and ensures all
|
|
required parameters are properly configured before proceeding.
|
|
|
|
Args:
|
|
input_data: Complete configuration for the batch prediction
|
|
Required keys:
|
|
- schedule_name (str): Schedule identifier for the prediction
|
|
- model_name (str): Name of the ML model to use
|
|
- model_id (int): Unique identifier for the model
|
|
- query (str): SQL query for data loading
|
|
- schema (dict, optional): Data schema definition
|
|
- table_name (str, optional): Target table for predictions
|
|
- input_filters (dict, optional): Data quality filters
|
|
- mlflow_transform_filters (dict, optional): MLFlow transform filters
|
|
- mlflow_predict_filters (dict, optional): MLFlow prediction filters
|
|
- model_retention (int, optional): Model retention period in minutes
|
|
- path_priority (list[str]): Decision path priority configuration
|
|
- opc_output_config (dict, optional): OPC server export configuration
|
|
- datetime_columns (list[str], optional): Columns to treat as datetime
|
|
|
|
Returns:
|
|
None: The workflow completes successfully when the child workflow finishes
|
|
|
|
Raises:
|
|
Exception: If any required parameters are missing or if the workflow fails
|
|
during data loading or workflow delegation
|
|
"""
|
|
|
|
metadata = {
|
|
'metadata': {
|
|
'schedule_name': input_data['schedule_name'],
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'workflow_name': 'predictions_batch',
|
|
}
|
|
}
|
|
|
|
# Load data using custom query
|
|
data = await workflow.execute_local_activity_method(
|
|
Activities.load_custom_query,
|
|
{
|
|
**metadata,
|
|
'query': input_data['query'],
|
|
'datetime_columns': input_data.get('datetime_columns', []),
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=300),
|
|
)
|
|
|
|
# Prepare input for prediction_process workflow
|
|
prediction_input = {
|
|
'metadata': metadata,
|
|
'data': data,
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
'transform_table_name': input_data['transform_table_name'],
|
|
'model_id': input_data['model_id'],
|
|
'model_name': input_data['model_name'],
|
|
'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
|
|
'mlflow_transform_filters': input_data.get(
|
|
'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
|
),
|
|
'mlflow_predict_filters': input_data.get(
|
|
'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
|
|
),
|
|
'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', {}),
|
|
'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}),
|
|
'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'),
|
|
'save_transform': input_data.get('save_transform', True),
|
|
}
|
|
|
|
# Execute prediction process workflow
|
|
await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input)
|