Files
sientia-dataops-model-manager/model_manager/workflows/predictions_batch.py
Bruno Domingues b102f79087 SIENTIAPDE-1243: Remove OPC server integration and add code quality tools.
This commit removes the OPC server integration from the Model Manager, including related activities, repositories, metrics, and configuration. It also adds code quality tools such as Ruff (linting/formatting), mypy (type checking), and Bandit (security analysis) along with a validation script and CI/CD integration for automated code validation. The README has been updated to reflect these changes.
2025-10-01 16:25:08 -03:00

126 lines
5.1 KiB
Python

from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from typing import Any
from sientia_do.temporal.policies import retry_policy
from datetime import timedelta
@workflow.defn(name="predictions_batch")
class PredictionsBatch():
"""
Main batch prediction workflow for the Model Manager 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
- 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'],
'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']),
'prediction_store_policy': input_data.get(
'prediction_store_policy', 'lts:1')
}
# Execute prediction process workflow
await workflow.execute_child_workflow(
'prediction_process', prediction_input)