sientia-dataops-model-manager/model-manager/workflows/predictions_batch.py

from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from laborious.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 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'],
            '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', {}),
            'prediction_store_policy': input_data.get(
                'prediction_store_policy', 'lts:1')
        }

        # Execute prediction process workflow
        await workflow.execute_child_workflow(
            'prediction_process', prediction_input)