vitor-aignosi
67fe4afaa6
Refactor and enhance the laborious workflow and utilities - Removed outdated test file `test_predictions_batch.py` from workflows. - Added `input_sample.json` for standardized input configuration. - Introduced `connectors_config.py` to manage database and service configurations. - Implemented a logging utility in `logger.py` for consistent logging across the application. - Created `policies.py` to define retry policies for workflows. - Developed comprehensive tests for `MLFlowRepository` in `test_model_repository.py`. - Added extensive tests for `OpcRepository` in `test_opc_repository.py`. - Updated `test_predictions_batch.py` to reflect new workflow structure and testing methodology.
90 lines
3.6 KiB
Python
90 lines
3.6 KiB
Python
from temporalio import workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from laborious.activities.activities import Activities
|
|
from typing import Any
|
|
from laborious.utils.policies import retry_policy
|
|
from datetime import timedelta
|
|
|
|
|
|
@workflow.defn(name="predictions_batch")
|
|
class PredictionsBatch():
|
|
@workflow.run
|
|
async def run(self, input_data: dict[str, Any]):
|
|
"""
|
|
This workflow runs a batch of predictions based on the input data.
|
|
|
|
The workflow executes in two main steps:
|
|
1. Prepares the activity with schedule and model information
|
|
2. Loads data using a custom query and executes the prediction process
|
|
|
|
Args:
|
|
input_data (dict[str, Any]): The input data for the workflow.
|
|
Contains the following keys:
|
|
schedule_name (str): The name of the schedule.
|
|
model_name (str): The name of the model.
|
|
model_id (int): The id of the model.
|
|
query (str): The SQL query to be executed to load data.
|
|
schema (dict, optional): The schema definition for the data.
|
|
table_name (str, optional): The name of the table to process.
|
|
input_filters (dict, optional): Filters to be applied during prediction.
|
|
mlflow_transform_filters (dict, optional): Filters to be applied during prediction.
|
|
mlflow_predict_filters (dict, optional): Filters to be applied during prediction.
|
|
model_retention (int, optional): The model retention period in minutes.
|
|
path_priority (list[str]): The path priority.
|
|
Returns:
|
|
None
|
|
|
|
Raises:
|
|
Exception: If any of the required parameters are missing or if the workflow fails.
|
|
"""
|
|
|
|
await workflow.execute_local_activity_method(
|
|
Activities.prepare_activity,
|
|
{
|
|
'schedule_name': input_data['schedule_name'],
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id'],
|
|
'workflow_name': 'predictions_batch'
|
|
},
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=60)
|
|
)
|
|
|
|
data = await workflow.execute_local_activity_method(
|
|
Activities.load_custom_query,
|
|
input_data['query'],
|
|
retry_policy=retry_policy,
|
|
start_to_close_timeout=timedelta(seconds=60)
|
|
)
|
|
|
|
# Prepare input for prediction_process workflow
|
|
prediction_input = {
|
|
'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_retention': input_data.get('model_retention', 60),
|
|
'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
|
|
'opc_output_config': input_data.get('opc_output_config', {})
|
|
}
|
|
|
|
await workflow.execute_child_workflow(
|
|
'prediction_process', prediction_input)
|