diff --git a/coverage.sh b/coverage.sh deleted file mode 100755 index 5867692..0000000 --- a/coverage.sh +++ /dev/null @@ -1 +0,0 @@ -pytest --cov=laborious --cov-report=html && xdg-open htmlcov/index.html \ No newline at end of file diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 919a2b2..605c20d 100644 --- a/laborious/activities/gates.py +++ b/laborious/activities/gates.py @@ -242,6 +242,28 @@ class Gates(BaseActivity): self.info("Nothing was filtered by the mlflow content gate", metadata) return None, 0, "" + def get_prediction_store_policy(self, + prediction_store_policy: str, + metadata: dict[str, Any]) -> tuple[str, int]: + policy_elements = prediction_store_policy.split(':') + + if len(policy_elements) < 2: + self.error( + f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata) + return 'lts', 1 + + policy_type = policy_elements[0] + policy_value = policy_elements[1] + + # If the policy_type is not lts or erl, we use the default policy + # If the policty_value is not a number or 0, we use the default policy + if policy_type not in ['lts', 'erl'] or not policy_value.isdigit() or int(policy_value) == 0: + self.error( + f"Invalid prediction store policy: {prediction_store_policy}, using default policy", metadata) + return 'lts', 1 + + return policy_type, int(policy_value) + @activity.defn(name="format_prediction") async def format_prediction(self, input_data: dict[str, Any]) -> dict[Any, Any]: """ @@ -252,21 +274,58 @@ class Gates(BaseActivity): - timestamp (str): The timestamp of the data. - model_id (str): The id of the model. - prediction_confidence (float): The confidence of the prediction. + - prediction_store_policy (str): The policy to store the prediction. Returns: dict: The formatted data. """ metadata = input_data['metadata'] + prediction_store_policy = input_data['prediction_store_policy'] self.info("Formatting prediction...", metadata) data = DataFrame(input_data['data']) - data['timestamp'] = input_data['timestamp'] + + self.debug( + f"Prediction store policy: {prediction_store_policy}", metadata) + + policy_type, policy_value = self.get_prediction_store_policy( + prediction_store_policy, metadata) + + # If data has no timestamp, we use the default timestamp and not sort the data + if 'timestamp' not in data.columns: + self.warning( + "Data has no timestamp, using default timestamp", metadata) + data['timestamp'] = input_data['timestamp'] + else: + self.debug( + f"Data has timestamp, sorting data by timestamp", metadata) + + # If policy_type is lts, we need to sort the data by timestamp descending and take the first policy_value rows + if policy_type == 'lts': + self.debug( + f"Sorting data by timestamp descending", metadata) + data = data.sort_values(by='timestamp', ascending=False) + # If policy_type is erl, we need to sort the data by timestamp ascending and take the first policy_value rows + elif policy_type == 'erl': + self.debug( + f"Sorting data by timestamp ascending", metadata) + data = data.sort_values(by='timestamp', ascending=True) + else: + self.error( + f"Invalid policy type: {policy_type}, using default policy", metadata) + raise ValueError( + f"Invalid policy type: {policy_type}") + + data = data.head(int(policy_value)) + data['model_id'] = input_data['model_id'] data['prediction_confidence'] = input_data['prediction_confidence'] data['prediction_status'] = 'Good' data['comments'] = "" - data = data.sort_values(by='timestamp') + data = data.sort_values(by='timestamp', ascending=False) + data = data.reset_index(drop=True) self.info(f"Prediction formatted: {data.size} rows", metadata) + self.debug(f"Prediction data: {data.to_string()}", metadata) return data.to_dict() diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index ee7c424..4e95c11 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -52,6 +52,8 @@ class FormatAndExportPrediction(): 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': prediction_confidence, + 'prediction_store_policy': input_data.get( + 'prediction_store_policy', 'lts:1') }, retry_policy=retry_policy, start_to_close_timeout=timedelta(seconds=60) diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 2032cc6..42b4d10 100644 --- a/tests/laborious/activities/test_gates.py +++ b/tests/laborious/activities/test_gates.py @@ -322,15 +322,68 @@ async def test_mlflow_content_gate_with_filter(gates_activity): gates_activity.send_notification.assert_called() +def test_get_prediction_store_policy_invalid_policy(gates_activity): + # Arrange + prediction_store_policy = 'INVALID_POLICY' + + # Act + policy_type, policy_value = gates_activity.get_prediction_store_policy( + prediction_store_policy, metadata) + + # Assert + assert policy_type == 'lts' + assert policy_value == 1 + + +def test_get_prediction_store_policy_invalid_policy_value(gates_activity): + # Arrange + prediction_store_policy = 'abc:INVALID_VALUE' + + # Act + policy_type, policy_value = gates_activity.get_prediction_store_policy( + prediction_store_policy, metadata) + + # Assert + assert policy_type == 'lts' + assert policy_value == 1 + + +def test_get_prediction_store_policy_valid_policy_type(gates_activity): + # Arrange + prediction_store_policy = 'abc:1' + + # Act + policy_type, policy_value = gates_activity.get_prediction_store_policy( + prediction_store_policy, metadata) + + # Assert + assert policy_type == 'lts' + assert policy_value == 1 + + +def test_get_prediction_store_policy_valid_policy(gates_activity): + # Arrange + prediction_store_policy = 'erl:1' + + # Act + policy_type, policy_value = gates_activity.get_prediction_store_policy( + prediction_store_policy, metadata) + + # Assert + assert policy_type == 'erl' + assert policy_value == 1 + + @mark.asyncio -async def test_format_prediction(gates_activity): +async def test_format_prediction_no_timestamp(gates_activity): # Arrange input_data = { **metadata, 'data': {'prediction': [1], 'response_time': [0.1]}, 'timestamp': '2023-05-26 11:12:27', 'model_id': 'test_model', - 'prediction_confidence': 0.9 + 'prediction_confidence': 0.9, + 'prediction_store_policy': 'lts:1' } # Act @@ -346,6 +399,83 @@ async def test_format_prediction(gates_activity): assert result['comments'] == {0: ""} +@mark.asyncio +async def test_format_prediction_with_timestamp_erl(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': {'prediction': [1, 2, 3], + 'response_time': [0.1, 0.2, 0.3], + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']}, + 'model_id': 'test_model', + 'prediction_confidence': 0.9, + 'prediction_store_policy': 'erl:2' + } + + # Act + result = await gates_activity.format_prediction(input_data) + + # Assert + assert result['prediction'] == {0: 2, 1: 1} + assert result['response_time'] == {0: 0.2, 1: 0.1} + assert result['timestamp'] == { + 0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'} + assert result['model_id'] == {0: 'test_model', 1: 'test_model'} + assert result['prediction_confidence'] == {0: 0.9, 1: 0.9} + assert result['prediction_status'] == {0: 'Good', 1: 'Good'} + assert result['comments'] == {0: "", 1: ""} + + +@mark.asyncio +async def test_format_prediction_with_timestamp_lts(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': {'prediction': [1, 2, 3], + 'response_time': [0.1, 0.2, 0.3], + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']}, + 'model_id': 'test_model', + 'prediction_confidence': 0.9, + 'prediction_store_policy': 'lts:2' + } + + # Act + result = await gates_activity.format_prediction(input_data) + + # Assert + assert result['prediction'] == {0: 3, 1: 2} + assert result['response_time'] == {0: 0.3, 1: 0.2} + assert result['timestamp'] == { + 0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'} + assert result['model_id'] == {0: 'test_model', 1: 'test_model'} + assert result['prediction_confidence'] == {0: 0.9, 1: 0.9} + assert result['prediction_status'] == {0: 'Good', 1: 'Good'} + assert result['comments'] == {0: "", 1: ""} + + +@mark.asyncio +async def test_format_prediction_with_timestamp_invalid_policy(gates_activity): + # Arrange + input_data = { + **metadata, + 'data': {'prediction': [1, 2, 3], + 'response_time': [0.1, 0.2, 0.3], + 'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29']}, + 'model_id': 'test_model', + 'prediction_confidence': 0.9, + 'prediction_store_policy': 'lts:2' + } + gates_activity.get_prediction_store_policy = MagicMock( + return_value=('invalid', 1)) + + try: + result = await gates_activity.format_prediction(input_data) + except ValueError as e: + assert str(e) == "Invalid policy type: invalid" + else: + assert False, "Expected ValueError" + + @mark.asyncio async def test_format_default_prediction(gates_activity): # Arrange diff --git a/tests/laborious/activities/test_mlflow.py b/tests/laborious/activities/test_mlflow.py index fc80f5a..7aced7c 100644 --- a/tests/laborious/activities/test_mlflow.py +++ b/tests/laborious/activities/test_mlflow.py @@ -82,7 +82,8 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow): } # Mock the transform response - expected_response = {'prediction': [0.5, 0.6]} + expected_response = {'prediction': [0.5, 0.6], 'timestamp': [ + '2024-01-01', '2024-01-02']} mlflow.model_monitoring_repository.transform.return_value = expected_response mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value @@ -98,7 +99,7 @@ async def test_request_transform(mock_max, mock_dataframe, mlflow): ) mock_dataframe = mock_dataframe.return_value.pivot.return_value mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True) - mock_dataframe.reset_index.assert_called_once() + # mock_dataframe.reset_index.assert_called_once() mock_dataframe.columns.name = None # Verify the response 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 54909a6..a8e6e20 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py @@ -35,7 +35,8 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): "schema": "test_schema", "table_name": "test_table", "opc_servers": ["test_server"], - "opc_output_config": {"test": "config"} + "opc_output_config": {"test": "config"}, + "prediction_store_policy": "erl:1" } await format_and_export_prediction.run(input_data) @@ -48,6 +49,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction): 'timestamp': input_data['timestamp'], 'model_id': input_data['model_id'], 'prediction_confidence': input_data['prediction_confidence'], + 'prediction_store_policy': input_data['prediction_store_policy'], **metadata }, retry_policy=ANY,