diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 56a1569..7c8ae76 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -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) diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 2c50abe..0447215 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -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() \ No newline at end of file + return reference_data \ No newline at end of file diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index c70cc81..9896fa7 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -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 diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index aa93ca5..404840e 100644 --- a/laborious/activities/opc.py +++ b/laborious/activities/opc.py @@ -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. diff --git a/laborious/utils/dataframe_utils.py b/laborious/utils/dataframe_utils.py new file mode 100644 index 0000000..c87764d --- /dev/null +++ b/laborious/utils/dataframe_utils.py @@ -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) + diff --git a/laborious/utils/temporal_codec.py b/laborious/utils/temporal_codec.py new file mode 100644 index 0000000..2c6d70a --- /dev/null +++ b/laborious/utils/temporal_codec.py @@ -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 + ) + diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index 1d28016..4700c48 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -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) diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py index 7afaef7..043ef7e 100644 --- a/laborious/workflows/drift.py +++ b/laborious/workflows/drift.py @@ -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}, diff --git a/laborious/workflows/minimal_retrain.py b/laborious/workflows/minimal_retrain.py index 02878cf..6ba88f9 100644 --- a/laborious/workflows/minimal_retrain.py +++ b/laborious/workflows/minimal_retrain.py @@ -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'], }, diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py index f5f02d2..f001962 100644 --- a/laborious/workflows/simple_metrics.py +++ b/laborious/workflows/simple_metrics.py @@ -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': { diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 4280ccf..1d59e7a 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -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, diff --git a/tests.ipynb b/tests.ipynb index 808277c..25f1cf0 100644 --- a/tests.ipynb +++ b/tests.ipynb @@ -588,6 +588,157 @@ "\n", "display(prediction_data)" ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "486b95b3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
ab
2025-01-0114
2025-01-0225
2025-01-0336
\n", + "
" + ], + "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": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
012
index2025-01-012025-01-022025-01-03
columnsabNone
data[1, 4][2, 5][3, 6]
\n", + "
" + ], + "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": {