vitor-aignosi
43f19ed93a
Update requirements.txt with new dependencies and refactor activity methods for improved functionality and error handling
86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
from temporalio import workflow
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
from laborious.activities.postgres import Postgres
|
|
from laborious.activities.mlflow import MLFlow
|
|
from laborious.activities.gates import Gates
|
|
from laborious.activities.opc import OPC
|
|
from typing import Any
|
|
|
|
|
|
@workflow.defn(name="predictions_batch")
|
|
class PredictionsBatch():
|
|
@workflow.run
|
|
async def run(self, input_data: dict[str, Any]):
|
|
|
|
await workflow.execute_activity_method(
|
|
Postgres.prepare_activity,
|
|
{
|
|
'schedule_name': input_data['schedule_name'],
|
|
'model_name': input_data['model_name'],
|
|
'model_id': input_data['model_id']
|
|
}
|
|
)
|
|
|
|
data = await workflow.execute_activity_method(
|
|
Postgres.load_custom_query,
|
|
input_data['query']
|
|
)
|
|
|
|
path_flag, confidence = await workflow.execute_activity_method(
|
|
Gates.input_gate,
|
|
{
|
|
'filters': input_data['filters'],
|
|
'data': data
|
|
}
|
|
)
|
|
|
|
if path_flag == 'stop':
|
|
return
|
|
|
|
if path_flag == 'continue':
|
|
# repeat last prediction
|
|
await workflow.execute_activity_method(
|
|
Postgres.repeat_last_prediction,
|
|
{
|
|
'schema': input_data['schema'],
|
|
'table_name': input_data['table_name'],
|
|
'model': input_data['model']
|
|
}
|
|
)
|
|
return
|
|
|
|
response_data, last_timestamp = await workflow.execute_activity_method(
|
|
MLFlow.transform_data,
|
|
{
|
|
'data': data,
|
|
'model_name': input_data['model_name'],
|
|
'model_retention': input_data['model_retention']
|
|
}
|
|
)
|
|
|
|
path_flag, confidence = await workflow.execute_activity_method(
|
|
Gates.mlflow_gate,
|
|
{
|
|
'filters': input_data['filters'],
|
|
'data': response_data,
|
|
'type': 'transform'
|
|
}
|
|
)
|
|
|
|
if path_flag == 'stop':
|
|
return
|
|
|
|
if path_flag is None:
|
|
# procced with prediction
|
|
response_data = await workflow.execute_activity_method(
|
|
MLFlow.request_predict,
|
|
{
|
|
'data': response_data,
|
|
'model_name': input_data['model_name'],
|
|
'model_retention': input_data['model_retention']
|
|
}
|
|
)
|
|
|
|
path_flag
|