diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index 58caf96..f15621a 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/e2e/test_predictions_batch_prediction_process.py @@ -3,7 +3,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios. """ from decimal import Decimal -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pandas as pd @@ -15,6 +15,7 @@ from temporalio.worker import Worker from e2e.helpers import ( assert_continue, assert_postgres_unique_violation_in_chain, + assert_prediction, assert_prediction_row_count, assert_repeat, assert_stop, @@ -25,6 +26,7 @@ from e2e.helpers import ( start_and_await_workflow, ) from laborious.activities.activities import Activities +from laborious.utils.models import minio_dataframe_payload as minio_payload_module from laborious.workflows.predictions_batch import PredictionsBatch DISTINCT_BATCH_TIMESTAMP = '2024-01-01 13:00:00+00:00' @@ -450,3 +452,37 @@ async def test_scenario_2_4_1_priority_conflict_resolution( prediction_confidence=Decimal(10), comments='Unknown MLFlow API error', ) + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_e2e_request_predict_inline_minio_payload_with_datetimeindex( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + mlflow_repository_stub, +): + """ + High offload threshold forces inline tabular dicts; ``DatetimeIndex`` must serialize as JSON + (string index keys via ``MinioDataFramePayload.from_dataframe``) so ``request_predict`` completes. + """ + client = temporal_test_env.client + model_id = 252 + with postgres_engine.begin() as conn: + conn.execute(text(f'DELETE FROM sientia_data.laborious_data WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM sientia_data.predictions WHERE model_id = {model_id}')) + conn.execute(text(f'DELETE FROM sientia_data.transformed_data WHERE model_id = {model_id}')) + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = get_base_input_data(model_id) + + with patch.object(minio_payload_module, 'OFFLOAD_THRESHOLD_BYTES', 10**9): + await start_and_await_workflow( + client, + PredictionsBatch.run, + input_data, + make_workflow_id('test-predict-inline-json-datetimeindex'), + ) + + assert_prediction(postgres_engine, model_id, prediction=0.5, prediction_confidence=0) + mlflow_repository_stub.stub_wrapper.predict.assert_called() diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 131193a..e5e1af5 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -324,8 +324,9 @@ class MLFlow(SientiaMonitoring): The activity normalizes ``NaN`` to ``None`` for JSON-friendly columns, sets the row index the same way as ``retrain_model`` (UTC ``DatetimeIndex`` from ``DATETIME_FORMAT_WITH_TZ``), - restores that index on the prediction frame, and records ``response_time``. Non-DataFrame - predictions are coerced to a single ``prediction`` column. + restores that index on the prediction frame, normalizes the prediction index to + ``DATETIME_FORMAT_WITH_TZ`` strings like ``request_transform``, and records ``response_time``. + Non-DataFrame predictions are coerced to a single ``prediction`` column. Args: - input_data: Same envelope as ``request_transform`` (``metadata``, ``model_name``, @@ -381,6 +382,7 @@ class MLFlow(SientiaMonitoring): predict_data.index = input_index predict_data['response_time'] = (end_time - start_time).total_seconds() + predict_data = self._detect_and_parse_datetime_index(predict_data, metadata) response_data: dict[str, Any] = {'success': True, 'content': predict_data} except Exception as e: diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index 30f792a..49e51c8 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -311,6 +311,7 @@ def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow): ) @patch('laborious.activities.mlflow.to_datetime') def test_request_predict_success_dataframe_and_meta(mock_to_datetime, mock_from_dataframe, mlflow): + mock_to_datetime.side_effect = lambda x, **kwargs: x wrapper = MagicMock() pred_df = pd.DataFrame({'raw': [0.3]}) wrapper.predict.return_value = (pred_df, {'m': 1})