SIENTIAPDE-1182

Implement prediction store policy handling in Gates activity

- Added a new method `get_prediction_store_policy` to validate and parse the prediction store policy.
- Updated `format_prediction` method to utilize the new policy handling, allowing for sorting of predictions based on the specified policy.
- Enhanced test coverage for the new policy handling, including various scenarios for valid and invalid policies.
- Removed the obsolete `coverage.sh` script.

coverage.sh

@@ -1 +0,0 @@
-pytest --cov=laborious --cov-report=html && xdg-open htmlcov/index.html

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()
 

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)

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

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

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,