gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-994

Browse Source 
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.
This commit is contained in:
vitor-aignosi
2025-05-23 17:34:47 -03:00
parent 5fe552410b
commit 67fe4afaa6
30 changed files with 1385 additions and 765 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
laborious/workflows/predictions_batch.py
View File

@@ -3,29 +3,87 @@ 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.
await workflow.execute_activity_method(
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_activity_method(
data = await workflow.execute_local_activity_method(
Activities.load_custom_query,
input_data['query']
input_data['query'],
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
input_data['data'] = data
# 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', input_data)
'prediction_process', prediction_input)

57
laborious/workflows/sub_workflows/format_and_export_prediction.py
View File

@@ -3,6 +3,8 @@ from temporalio import workflow
with workflow.unsafe.imports_passed_through():
from laborious.activities.activities import Activities
from typing import Any
from datetime import timedelta
from laborious.utils.policies import retry_policy
@workflow.defn(name="format_and_export_prediction")
@@ -11,23 +13,25 @@ class FormatAndExportPrediction():
async def run(self, input_data: dict[str, Any]):
"""
This workflow formats and exports predictions based on path_flag:
- If path_flag is None: formats prediction using input data, timestamp, model_id and confidence
- If path_flag exists: creates default prediction with timestamp, model_id, confidence and comment
- If path_flag is None: formats prediction
using input data, timestamp, model_id and confidence
- If path_flag exists: creates default prediction
with timestamp, model_id, confidence and comment
Finally exports formatted prediction to postgres table
Args:
input_data(dict[str, Any]): The input data for the workflow. Contains the following keys:
- path_flag(str): The path flag to determine the type of prediction to format
- data(dict[str, Any]): The data to format
- prediction_confidence(float): The prediction confidence to be registered
- timestamp(str): The timestamp of the prediction, synchronized with the data
- model_id(str): The model id of the prediction
- model_name(str): The model name of the prediction
- model_retention(str): The model retention of the prediction
- comment(str): The comment to be registered
- schema(str): The schema of the prediction
- table_name(str): The table name of the prediction
- opc_servers(list[str]): The opc servers of the prediction
- opc_output_config(dict[str, Any]): The opc output config of the prediction
input_data(dict[str, Any]): The input data for the workflow.
Contains the following keys:
path_flag(str): The path flag to determine the type of prediction to format
data(dict[str, Any]): The data to format
prediction_confidence(float): The prediction confidence to be registered
timestamp(str): The timestamp of the prediction, synchronized with the data
model_id(int): The model id of the prediction
model_name(str): The model name of the prediction
model_retention(str): The model retention of the prediction
comment(str): The comment to be registered
schema(str): The schema of the prediction
table_name(str): The table name of the prediction
opc_output_config(dict[str, Any]): The opc output config of the prediction
Returns:
bool: True if the workflow was successful, False otherwise.
@@ -36,28 +40,34 @@ class FormatAndExportPrediction():
data = input_data['data']
prediction_confidence = input_data['prediction_confidence']
print(f"Input data: {input_data}")
if path_flag is None:
# proceed with formatting and exporting
prediction = await workflow.execute_activity_method(
prediction = await workflow.execute_local_activity_method(
Activities.format_prediction,
{
'data': data,
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
else:
# create default prediction
prediction = await workflow.execute_activity_method(
prediction = await workflow.execute_local_activity_method(
Activities.format_default_prediction,
{
'timestamp': input_data['timestamp'],
'model_id': input_data['model_id'],
'prediction_confidence': prediction_confidence,
'comment': input_data['comment']
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
# write to postgres
@@ -67,17 +77,20 @@ class FormatAndExportPrediction():
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
# write to opc
opc_holder = workflow.execute_activity_method(
Activities.write_opc_data,
{
'opc_servers': input_data['opc_servers'],
'opc_output_config': input_data['opc_output_config'],
'data': prediction
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(seconds=60)
)
await postgres_holder

159
laborious/workflows/sub_workflows/prediction_process.py
View File

@@ -3,101 +3,144 @@ 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="prediction_process")
class PredictionProcess():
@workflow.run
async def run(self, input_data: dict[str, Any]):
"""
This workflow runs a prediction process 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:
data (dict[str, Any]): The data to be used for the prediction.
schema (str): The schema of the table.
table_name (str): The name of the table.
model_id (int): The id of the model.
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_name (str): The name of the model.
model_retention (int, optional): The model retention period in minutes.
path_priority (list[str]): The path priority.
opc_output_config (dict[str, Any]): The opc output config of the prediction.
Returns:
None
Raises:
Exception: If any of the required parameters are missing or if the workflow fails.
"""
data = input_data['data']
schema = input_data['schema']
table_name = input_data['table_name']
model = input_data['model']
filters = input_data['filters']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
last_timestamp = await workflow.execute_activity_method(
last_timestamp = await workflow.execute_local_activity_method(
Activities.get_last_timestamp,
{
'data': data
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence = await workflow.execute_activity_method(
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.input_gate,
{
'filters': input_data['filters'],
'data': data
}
'filters': input_data['input_filters'],
'data': data,
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
response_data = await workflow.execute_activity_method(
response_data = await workflow.execute_local_activity_method(
Activities.request_transform,
{
'data': data,
'model_name': model_name,
'model_retention': model_retention
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence = await workflow.execute_activity_method(
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
'filters': filters,
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform'
}
'type': 'transform',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
path_flag, confidence = await workflow.execute_activity_method(
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_content_gate,
{
'filters': filters,
'filters': input_data['mlflow_transform_filters'],
'data': response_data,
'type': 'transform'
}
'type': 'transform',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
response_data = await workflow.execute_activity_method(
response_data = await workflow.execute_local_activity_method(
Activities.request_predict,
{
'data': response_data,
'model_name': model_name,
'model_retention': model_retention
}
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
path_flag, confidence = await workflow.execute_activity_method(
path_flag, confidence, comment = await workflow.execute_local_activity_method(
Activities.mlflow_response_gate,
{
'filters': filters,
'filters': input_data['mlflow_predict_filters'],
'data': response_data,
'type': 'predict'
}
'type': 'predict',
'path_priority': input_data['path_priority']
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
if await self.path_flag_handler(
data, path_flag, confidence, schema, table_name,
model, last_timestamp, model_name, model_retention
data, path_flag, input_data, confidence, last_timestamp, comment
):
return
@@ -108,60 +151,78 @@ class PredictionProcess():
'data': response_data['content'],
'prediction_confidence': confidence,
'timestamp': response_data['timestamp'],
'model_id': model,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention
'model_retention': model_retention,
'opc_output_config': input_data['opc_output_config']
}
)
async def path_flag_handler(self, data: dict[str, Any], path_flag: str,
confidence: int, schema: str, table_name: str,
model: str, last_timestamp: str, model_name: str,
model_retention: str):
input_data: dict[str, Any], confidence: int,
last_timestamp: str, comment: str):
"""
This function handles the path flag and the confidence of the prediction.
It returns True if the prediction should be stopped. If path_flag is 'repeat', it repeats the last prediction.
If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop', it stops the prediction process.
It returns True if the prediction should be stopped. If path_flag is 'repeat',
it repeats the last prediction.
If path_flag is 'continue', it calls the write workflow. If path_flag is 'stop',
it stops the prediction process.
Args:
data (dict[str, Any]): The data to be used for the prediction.
path_flag (str): The path flag to determine the type of prediction to format
confidence (int): The confidence of the prediction
schema (str): The schema of the prediction
table_name (str): The table name of the prediction
model (str): The model id of the prediction
model_id (int): The model id of the prediction
last_timestamp (str): The timestamp of the last prediction
model_name (str): The model name of the prediction
model_retention (str): The model retention of the prediction
model_retention (int): The model retention of the prediction
comment (str): The comment of the prediction
Returns:
bool: True if the prediction should be stopped, False otherwise.
"""
if path_flag == 'stop':
schema = input_data['schema']
table_name = input_data['table_name']
model_id = input_data['model_id']
model_name = input_data['model_name']
model_retention = input_data['model_retention']
path_flag = path_flag.upper() if path_flag else None
if path_flag == 'STOP':
return True
elif path_flag == 'repeat':
elif path_flag == 'REPEAT':
# repeat last prediction
await workflow.execute_activity_method(
Activities.repeat_last_prediction,
{
'schema': schema,
'table_name': table_name,
'model': model
}
'model_id': model_id
},
retry_policy=retry_policy,
start_to_close_timeout=timedelta(minutes=1),
)
return True
elif path_flag == 'continue':
elif path_flag == 'CONTINUE':
# call write workflow
workflow.execute_child_workflow(
await workflow.execute_child_workflow(
'format_and_export_prediction',
{
'path_flag': path_flag,
'data': data,
'prediction_confidence': confidence,
'timestamp': last_timestamp,
'model_id': model,
'model_id': model_id,
'model_name': model_name,
'model_retention': model_retention
'model_retention': model_retention,
'schema': schema,
'table_name': table_name,
'comment': comment,
'opc_output_config': input_data['opc_output_config']
}
)
return True