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

SIENTIAPDE-1273

Browse Source 
Refactor data handling in various modules to ensure DataFrame consistency

- Replaced direct DataFrame instantiation with `ensure_dataframe` utility in Gates, MLFlow, OPC, and ModelMetrics classes to standardize data handling.
- Updated return types in several asynchronous methods to return DataFrames instead of dictionaries for improved usability.
- Adjusted data export processes in workflows to convert DataFrames to dictionaries with `to_dict(orient='records')` for compatibility with downstream systems.
This commit is contained in:
vitor-aignosi
2025-11-14 16:55:00 -03:00
parent 4191f68d3b
commit d2b365a34d
12 changed files with 368 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
laborious/activities/gates.py
View File

@@ -15,6 +15,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
from laborious import metrics
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.filters.conditional_filters import (
filter_empty_data,
filter_specific_variables_null_values,
@@ -149,7 +150,7 @@ class Gates(SientiaMonitoring):
self.info('Performing input gate...', metadata)
filters = input_data['filters']
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
path_priority = input_data['path_priority']
filter_output = []
@@ -404,7 +405,7 @@ class Gates(SientiaMonitoring):
@activity.defn(name='format_transformed_data')
async def format_transformed_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
async def format_transformed_data(self, input_data: dict[str, Any]) -> DataFrame:
"""
Format transformed data according to configured storage policies.
"""
@@ -414,7 +415,7 @@ class Gates(SientiaMonitoring):
self.info('Formatting transformed data...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
data['timestamp'] = data.index
data = data.reset_index(drop=True)
@@ -422,10 +423,10 @@ class Gates(SientiaMonitoring):
data = data.melt(id_vars='timestamp', var_name='variable', value_name='value')
data['model_id'] = model_id
return data.to_dict()
return data
@activity.defn(name='format_prediction')
async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
async def format_prediction(self, input_data: dict[str, Any]) -> DataFrame:
"""
Format prediction data according to configured storage policies.
@@ -453,7 +454,7 @@ class Gates(SientiaMonitoring):
prediction_store_policy = input_data['prediction_store_policy']
self.info('Formatting prediction...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
# Create timestamp column from index and reset index
data['timestamp'] = data.index
@@ -495,10 +496,10 @@ class Gates(SientiaMonitoring):
self.info(f'Prediction formatted: {len(data)} rows', metadata)
self.debug(f'Prediction data: {data.head(5).to_string()}', metadata)
return data.to_dict()
return data
@activity.defn(name='format_default_prediction')
async def format_default_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]:
async def format_default_prediction(self, input_data: dict[str, Any]) -> DataFrame:
"""
Create and format default prediction data for error conditions.
@@ -540,10 +541,10 @@ class Gates(SientiaMonitoring):
)
self.info(f'Default prediction formatted: {data.size} rows', metadata)
return data.to_dict()
return data
@activity.defn(name='format_retrain_report')
async def format_retrain_report(self, input_data: dict[str, Any]) -> dict[Any, Any]:
async def format_retrain_report(self, input_data: dict[str, Any]) -> DataFrame:
"""
Format retrain report data according to configured storage policies.
"""
@@ -572,7 +573,7 @@ class Gates(SientiaMonitoring):
self.debug(f'Retrain report: {report.to_csv()}', metadata)
return report.to_dict()
return report
@activity.defn(name='get_last_timestamp')
async def get_last_timestamp(self, input_data: dict[str, Any]) -> str:
@@ -601,7 +602,7 @@ class Gates(SientiaMonitoring):
self.info('Getting last timestamp...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
self.debug(f'Input data: {data.head(5).to_string()}', metadata)

9
laborious/activities/mlflow.py
View File

@@ -20,6 +20,7 @@ with workflow.unsafe.imports_passed_through():
now,
)
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.repository.minio_repository import MinioRepository
from laborious.utils.repository.model_repository import MLFlowRepository
@@ -139,7 +140,7 @@ class MLFlow(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Transforming data...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
@@ -211,7 +212,7 @@ class MLFlow(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Predicting data...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
model_name = input_data['model_name']
model_config = input_data.get('model_config', {})
@@ -422,7 +423,7 @@ class MLFlow(SientiaMonitoring):
@activity.defn(name='get_reference_data')
async def get_reference_data(self, input_data: dict[str, Any]) -> dict[Hashable, Any] | None:
async def get_reference_data(self, input_data: dict[str, Any]) -> DataFrame | None:
"""
Get reference data from the MLflow Model Registry.
@@ -451,4 +452,4 @@ class MLFlow(SientiaMonitoring):
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
return reference_data.to_dict()
return reference_data

8
laborious/activities/model_metrics.py
View File

@@ -138,7 +138,7 @@ class ModelMetrics(SientiaMonitoring):
@activity.defn(name='calculate_drift')
async def calculate_drift(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
async def calculate_drift(self, input_data: dict[str, Any]) -> DataFrame | dict:
"""
Calculate drift metrics for a model.
@@ -262,9 +262,9 @@ class ModelMetrics(SientiaMonitoring):
self.debug(f'Drift dataframe: Size {drift_df.shape} \n{drift_df.head(5).to_string()}', metadata)
return drift_df.to_dict()
return drift_df
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> dict[Hashable, Any]:
async def calculate_simple_metrics(self, input_data: dict[str, Any]) -> DataFrame:
"""
Calculate simple metrics for a model. Metrics available are:
- rmse
@@ -342,6 +342,6 @@ class ModelMetrics(SientiaMonitoring):
self.debug(f'Simple metrics dataframe: Size {data.shape} \n{data.head(5).to_string()}', metadata)
return data.to_dict()
return data

8
laborious/activities/opc.py
View File

@@ -1,3 +1,4 @@
from typing import Hashable
from temporalio import activity, workflow
with workflow.unsafe.imports_passed_through():
@@ -11,6 +12,7 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from laborious.utils.dataframe_utils import ensure_dataframe
from laborious.utils.repository.opc_repository import OpcRepository
OPC_WRITTING_ERROR_CONFIDENCE = 12
@@ -275,7 +277,7 @@ class OPC(SientiaMonitoring):
@activity.defn(name='write_opc_data')
async def write_opc_data(
self, input_data: dict[str, Any]
) -> tuple[dict[Any, Any], dict[str, dict[str, float | None]]]:
) -> tuple[dict[Hashable, Any], dict[str, dict[str, float | None]]]:
"""
Write prediction and confidence data to OPC servers. The two writing
operations are optional and independent of each other.
@@ -295,7 +297,7 @@ class OPC(SientiaMonitoring):
"""
metadata = input_data['metadata']
self.info('Writing data to OPC servers...', metadata)
data = DataFrame(input_data['data'])
data = ensure_dataframe(input_data['data'])
opc_output_config = input_data['opc_output_config']
self.info(f'Data to write: {data.size} rows', metadata)
@@ -324,7 +326,7 @@ class OPC(SientiaMonitoring):
def process_confidence(
self, data: DataFrame, success: bool, metadata: dict[str, Any]
) -> dict[Any, Any]:
) -> dict[Hashable, Any]:
"""
Process prediction confidence based on OPC write operation success.

31
laborious/utils/dataframe_utils.py Normal file
View File

@@ -0,0 +1,31 @@
"""
DataFrame utility functions for handling serialized DataFrames.
This module provides helper functions to work with DataFrames that may
come from Temporal serialization (already as DataFrame) or from legacy
code (as dict).
"""
from typing import Any
from pandas import DataFrame
def ensure_dataframe(data: Any) -> DataFrame:
"""
Ensure that data is a DataFrame, converting from dict if necessary.
This function handles both cases:
- Data already deserialized as DataFrame (from Temporal codec)
- Data as dict (legacy format or non-DataFrame serialization)
Args:
data: Data that should be a DataFrame (can be DataFrame or dict)
Returns:
DataFrame: The data as a pandas DataFrame
"""
if isinstance(data, DataFrame):
return data
return DataFrame(data)

149
laborious/utils/temporal_codec.py Normal file
View File

@@ -0,0 +1,149 @@
"""
Temporal Codec for DataFrame Serialization
This module provides a custom Temporal DataConverter that automatically
serializes pandas DataFrames to Parquet format and deserializes them back.
The codec only handles DataFrames, leaving all other types to the default
Temporal serialization mechanism.
"""
import io
from collections.abc import Sequence
from typing import Any, List, Optional, Type
from temporalio.api.common.v1 import Payload
from temporalio.converter import DataConverter, PayloadConverter
from pandas import DataFrame, read_parquet
class DataFramePayloadConverter(PayloadConverter):
"""
Custom PayloadConverter that serializes pandas DataFrames to Parquet format.
Only DataFrames are handled by this converter. All other types are passed
to the default Temporal serialization mechanism.
"""
def __init__(self, default_payload_converter: PayloadConverter):
"""
Initialize the DataFrame payload converter.
Args:
default_payload_converter: The default Temporal payload converter
to use for non-DataFrame types
"""
self._default = default_payload_converter
def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
"""
Convert values to Temporal Payloads.
If a value is a pandas DataFrame, it is serialized to Parquet format.
Otherwise, the default converter is used.
Args:
values: The values to serialize
Returns:
List[Payload]: The serialized payloads
"""
payloads = []
for value in values:
# Check if value is a DataFrame
try:
if isinstance(value, DataFrame):
buffer = io.BytesIO()
value.to_parquet(buffer, engine='pyarrow', index=True)
payloads.append(Payload(
metadata={"encoding": b"parquet-dataframe"},
data=buffer.getvalue()
))
continue
except ImportError:
# pandas not available, fall through to default
pass
except Exception:
# Error serializing DataFrame, fall through to default
pass
# Use default converter for all other types
default_payloads = self._default.to_payloads([value])
payloads.extend(default_payloads)
return payloads
def from_payloads(
self,
payloads: Sequence[Payload],
type_hints: Optional[List[Type]] = None,
) -> List[Any]:
"""
Convert Temporal Payloads back to Python values.
If a payload metadata indicates it's a Parquet-serialized DataFrame,
it is deserialized. Otherwise, the default converter is used.
Args:
payloads: The payloads to deserialize
type_hints: Optional type hints for the expected return types
Returns:
List[Any]: The deserialized values
"""
values = []
for i, payload in enumerate(payloads):
# Check if this is a Parquet-serialized DataFrame
if payload.metadata.get("encoding") == b"parquet-dataframe":
try:
buffer = io.BytesIO(payload.data)
values.append(read_parquet(buffer))
continue
except ImportError:
# pandas not available, fall through to default
pass
except Exception:
# Error deserializing DataFrame, fall through to default
pass
# Use default converter for all other types
# type_hints must have same length as payloads if provided
type_hint = type_hints[i] if type_hints and i < len(type_hints) else None
default_values = self._default.from_payloads([payload], [type_hint] if type_hint is not None else None)
values.extend(default_values)
return values
def create_dataframe_data_converter() -> DataConverter:
"""
Create a DataConverter with DataFrame serialization support.
Returns:
DataConverter: A DataConverter that handles DataFrames automatically
"""
# Get default converter to use as fallback
# DataConverter.default is an attribute, not a method
default_converter = DataConverter.default
# Create a factory class that extends PayloadConverter
class DataFramePayloadConverterFactory(PayloadConverter):
def __init__(self):
super().__init__()
# Create instance of default payload converter to use as fallback
self._default_converter = default_converter.payload_converter_class()
def to_payloads(self, values: Sequence[Any]) -> List[Payload]:
return DataFramePayloadConverter(self._default_converter).to_payloads(values)
def from_payloads(
self,
payloads: Sequence[Payload],
type_hints: Optional[List[Type]] = None,
) -> List[Any]:
return DataFramePayloadConverter(self._default_converter).from_payloads(payloads, type_hints)
return DataConverter(
payload_converter_class=DataFramePayloadConverterFactory
)

5
laborious/worker/worker.py
View File

@@ -48,6 +48,7 @@ with workflow.unsafe.imports_passed_through():
build_opc_config,
build_postgres_config,
)
from laborious.utils.temporal_codec import create_dataframe_data_converter
from laborious.workflows.minimal_retrain import MinimalRetrain
from laborious.workflows.predictions_batch import PredictionsBatch
from laborious.workflows.drift import Drift
@@ -132,10 +133,14 @@ async def main():
logger.custom_info(f'Starting Temporal Client at {host}...', metadata)
# Create custom data converter with DataFrame support
data_converter = create_dataframe_data_converter()
temporal_client = await client.Client.connect(
target_host=host,
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
runtime=new_runtime,
data_converter=data_converter,
)
logger.custom_info('Starting Workers...', metadata)

2
laborious/workflows/drift.py
View File

@@ -93,7 +93,7 @@ class Drift:
Activities.export_data_to_postgres,
{
**metadata,
'data': drift_data,
'data': drift_data.to_dict(orient='records'),
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},

2
laborious/workflows/minimal_retrain.py
View File

@@ -127,7 +127,7 @@ class MinimalRetrain:
Activities.export_data_to_postgres,
{
**metadata,
'data': report,
'data': report.to_dict(orient='records'),
'schema': input_data['schema'],
'table_name': input_data['table_name'],
},

2
laborious/workflows/simple_metrics.py
View File

@@ -82,7 +82,7 @@ class SimpleMetrics:
Activities.export_data_to_postgres,
{
**metadata,
'data': simple_metrics,
'data': simple_metrics.to_dict(orient='records'),
'schema': input_data['schema'],
'table_name': input_data['target_table_name'],
'timestamp_conversion': {

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

@@ -107,7 +107,7 @@ class FormatAndExportPrediction:
**metadata,
'schema': input_data['schema'],
'table_name': input_data['transform_table_name'],
'data': transformed,
'data': transformed.to_dict(orient='records'),
'timestamp_conversion': {
'column': 'timestamp',
'format': DATETIME_FORMAT_WITH_TZ,
@@ -156,7 +156,7 @@ class FormatAndExportPrediction:
**metadata,
'schema': input_data['schema'],
'table_name': input_data['table_name'],
'data': prediction,
'data': prediction.to_dict(orient='records'),
'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
},
retry_policy=retry_policy,

151
tests.ipynb
View File

@@ -588,6 +588,157 @@
"\n",
"display(prediction_data)"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "486b95b3",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>a</th>\n",
" <th>b</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>2025-01-01</th>\n",
" <td>1</td>\n",
" <td>4</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2025-01-02</th>\n",
" <td>2</td>\n",
" <td>5</td>\n",
" </tr>\n",
" <tr>\n",
" <th>2025-01-03</th>\n",
" <td>3</td>\n",
" <td>6</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" a b\n",
"2025-01-01 1 4\n",
"2025-01-02 2 5\n",
"2025-01-03 3 6"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/plain": [
"{'index': ['2025-01-01', '2025-01-02', '2025-01-03'],\n",
" 'columns': ['a', 'b'],\n",
" 'data': [[1, 4], [2, 5], [3, 6]]}"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>0</th>\n",
" <th>1</th>\n",
" <th>2</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>index</th>\n",
" <td>2025-01-01</td>\n",
" <td>2025-01-02</td>\n",
" <td>2025-01-03</td>\n",
" </tr>\n",
" <tr>\n",
" <th>columns</th>\n",
" <td>a</td>\n",
" <td>b</td>\n",
" <td>None</td>\n",
" </tr>\n",
" <tr>\n",
" <th>data</th>\n",
" <td>[1, 4]</td>\n",
" <td>[2, 5]</td>\n",
" <td>[3, 6]</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" 0 1 2\n",
"index 2025-01-01 2025-01-02 2025-01-03\n",
"columns a b None\n",
"data [1, 4] [2, 5] [3, 6]"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from pandas import DataFrame\n",
"\n",
"data = DataFrame({\n",
" \"a\": {\"2025-01-01\": 1, \"2025-01-02\": 2, \"2025-01-03\": 3},\n",
" \"b\": {\"2025-01-01\": 4, \"2025-01-02\": 5, \"2025-01-03\": 6},\n",
"})\n",
"\n",
"display(data)\n",
"\n",
"data_list = data.to_dict('split')\n",
"\n",
"display(data_list)\n",
"\n",
"data_rec = DataFrame.from_dict(data_list, orient='index')\n",
"\n",
"display(data_rec)"
]
}
],
"metadata": {