diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index b53562e..ad88ca1 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -522,7 +522,7 @@ class Gates(MinioManager): operation='transform', workflow_metadata=metadata, last_timestamp=payload.last_timestamp, - logger=self.logger + logger=self.logger, ) @activity.defn(name='format_prediction') diff --git a/laborious/activities/mlflow.py b/laborious/activities/mlflow.py index 207e20e..f4d42a3 100644 --- a/laborious/activities/mlflow.py +++ b/laborious/activities/mlflow.py @@ -19,8 +19,8 @@ with workflow.unsafe.imports_passed_through(): ) from sientia_do.utils.formatters import create_sample_dict - from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload from laborious.utils.dataframe_debug import build_dataframe_debug_message + from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload from laborious.utils.repository.minio_manager import MinioManager from laborious.utils.repository.model_repository import MLFlowRepository @@ -43,6 +43,7 @@ class MLFlow(MinioManager): mlflow_password (str): MLFlow authentication password model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 def __init__( diff --git a/laborious/activities/model_metrics.py b/laborious/activities/model_metrics.py index b336fbd..79b140c 100644 --- a/laborious/activities/model_metrics.py +++ b/laborious/activities/model_metrics.py @@ -27,13 +27,13 @@ warnings.filterwarnings( class ModelMetrics(SientiaMonitoring): """ - _MAX_DEBUG_DATAFRAME_ROWS = 100 - Metrics activities for the Laborious system. This class provides activities for writing metrics to the Prometheus monitoring system. """ + _MAX_DEBUG_DATAFRAME_ROWS = 100 + def __init__( self, logger: Logger, @@ -99,7 +99,9 @@ class ModelMetrics(SientiaMonitoring): model_analysis = ModelAnalysis(config=config) - self._debug_dataframe(f'Reference data: Size {reference_data.shape}', reference_data, metadata) + self._debug_dataframe( + f'Reference data: Size {reference_data.shape}', reference_data, metadata + ) self._debug_dataframe(f'Target data: Size {target_data.shape}', target_data, metadata) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 581e57f..b4fdad5 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -474,7 +474,9 @@ class MLFlowRepository(SientiaMonitoring): raw_model = mlflow.pyfunc.load_model(artifact_path) model = raw_model._model_impl.python_model - self.debug(f"Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}", metadata) + self.debug( + f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata + ) else: if model_type == 'predict': model = await self.load_predict_model(model_name, metadata, flavor) @@ -1293,7 +1295,9 @@ class MLFlowRepository(SientiaMonitoring): end_time = datetime.now() if isinstance(predict_data, pd.DataFrame): - self._debug_dataframe('Data received from model prediction:', predict_data, metadata) + self._debug_dataframe( + 'Data received from model prediction:', predict_data, metadata + ) # predict_data.to_csv( # f"tmp/predicted_data_{model_name}.csv", index=True) diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 3d072b9..c144068 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -541,6 +541,7 @@ async def test_format_prediction_no_timestamp(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:1', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -580,6 +581,7 @@ async def test_format_prediction_with_timestamp_erl(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'erl:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -619,6 +621,7 @@ async def test_format_prediction_with_timestamp_lts(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', + 'timestamp': '2023-05-26 11:12:27', } # Act @@ -655,6 +658,7 @@ async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): 'model_id': 'test_model', 'prediction_confidence': 0.9, 'prediction_store_policy': 'lts:2', + 'timestamp': '2023-05-26 11:12:27', } gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1)) diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index daaebab..5d242e5 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -181,6 +181,7 @@ async def test_request_transform_failure(mock_from_dataframe, mlflow): status=transform_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, + logger=mlflow.logger, ) assert response_data == mock_from_dataframe.return_value @@ -252,6 +253,7 @@ async def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, ml status=predict_response, workflow_metadata=metadata['metadata'], last_timestamp=payload.last_timestamp, + logger=mlflow.logger, ) assert response_data == mock_from_dataframe.return_value diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index 6e8817f..4fbc62b 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -93,6 +93,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -100,7 +101,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -243,6 +245,7 @@ async def test_run_none_path_flag_with_transformed_data( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -250,7 +253,8 @@ async def test_run_none_path_flag_with_transformed_data( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -349,14 +353,16 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, - **metadata, 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -456,6 +462,7 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -463,7 +470,8 @@ async def test_run_none_path_flag_with_pi_web_api(workflow_mock, format_and_expo 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -560,6 +568,7 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': prediction_data, @@ -567,7 +576,8 @@ async def test_run_none_path_flag_with_pi_web_api_and_opc( 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, @@ -648,6 +658,7 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e call( Activities.export_data_to_postgres, { + **metadata, 'schema': input_data['schema'], 'table_name': input_data['table_name'], 'data': pi_web_api_data, @@ -655,7 +666,8 @@ async def test_run_default_path_flag_with_pi_web_api(workflow_mock, format_and_e 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, - **metadata, + 'on_conflict': 'error', + 'unique_columns': ['model_id', 'timestamp'], }, retry_policy=ANY, start_to_close_timeout=ANY, diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 60f3744..a018504 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -51,14 +51,17 @@ async def test_run(workflow_mock, prediction_process): } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict # mlflow_response_gate (predict) ('continue', 0.95, 'Error'), ] @@ -67,7 +70,7 @@ async def test_run(workflow_mock, prediction_process): await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 6 + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -83,7 +86,7 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -130,7 +133,7 @@ async def test_run(workflow_mock, prediction_process): ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, @@ -166,6 +169,7 @@ async def test_run(workflow_mock, prediction_process): 'subworkflow.format_and_export_prediction', { 'metadata': metadata, + 'on_conflict': 'error', 'path_flag': 'continue', 'data': {'content': 'predicted_data', 'timestamp': '2024-01-01'}, 'transformed_data': {'content': 'transformed_data', 'timestamp': '2024-01-01'}, @@ -266,9 +270,12 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('repeat', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data ('continue', 0.95, 'Error'), # mlflow_response_gate (transform) ] @@ -276,7 +283,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 3 + assert workflow_mock.execute_local_activity_method.call_count == 2 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -292,7 +299,7 @@ async def test_run_stop_at_first_mlflow_response_gate(workflow_mock, prediction_ ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -353,9 +360,12 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) @@ -366,7 +376,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 4 + assert workflow_mock.execute_local_activity_method.call_count == 3 workflow_mock.execute_local_activity_method.assert_has_calls( [ @@ -383,7 +393,7 @@ async def test_run_stop_at_mlflow_content_gate(workflow_mock, prediction_process ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -460,14 +470,17 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p } # Mock the activity responses + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'Input data with bad quality'), # input_gate - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, # transform_data # mlflow_response_gate (transform) ('continue', 0.95, 'Error'), # mlflow_content_gate (transform) ('continue', 0.95, 'Transformed data not passed the content filter'), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, # request_predict ('continue', 0.95, 'Error'), # mlflow_response_gate (predict) ] @@ -475,7 +488,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p await prediction_process.run(input_data) # Assert - assert workflow_mock.execute_local_activity_method.call_count == 6 + assert workflow_mock.execute_local_activity_method.call_count == 4 workflow_mock.execute_local_activity_method.assert_has_calls( [ call( @@ -491,7 +504,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_transform, @@ -538,7 +551,7 @@ async def test_run_stop_at_mlflow_last_response_gate(workflow_mock, prediction_p ) ] ) - workflow_mock.execute_local_activity_method.assert_has_calls( + workflow_mock.execute_activity_method.assert_has_calls( [ call( Activities.request_predict, @@ -727,6 +740,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): 'confidence_tags': {}, }, 'prediction_store_policy': prediction_store_policy, + 'on_conflict': 'error', }, ) @@ -806,12 +820,15 @@ async def test_run_with_cleanup_prefixes(workflow_mock, prediction_process): 'prediction_store_policy': 'lts:1', } + workflow_mock.execute_activity_method.side_effect = [ + {'content': 'transformed_data', 'timestamp': '2024-01-01'}, + {'content': 'predicted_data', 'timestamp': '2024-01-01'}, + MagicMock(), + ] workflow_mock.execute_local_activity_method.side_effect = [ ('continue', 0.95, 'ok'), - {'content': 'transformed_data', 'timestamp': '2024-01-01'}, ('continue', 0.95, ''), ('continue', 0.95, ''), - {'content': 'predicted_data', 'timestamp': '2024-01-01'}, ('continue', 0.95, ''), ] diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index de52ad8..14280d6 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -95,6 +95,7 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch 'model_config': input_data.get('model_config', {}), 'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']), 'opc_output_config': input_data.get('opc_output_config', {}), + 'on_conflict': input_data.get('on_conflict', 'error'), 'pi_web_api_output_config': input_data.get('pi_web_api_output_config', {}), 'prediction_store_policy': input_data.get('prediction_store_policy', 'lts:1'), 'save_transform': input_data.get('save_transform', True),