SIENTIAPDE-1325

Enhance metrics handling in Gates and MLFlow classes

- Added checks for `None` response times before emitting OPC writing metrics in the Gates class to prevent unnecessary metric emissions.
- Updated the MLFlow class to conditionally sort and drop duplicates based on the presence of the 'created_at' column, ensuring robustness in data processing.
- Adjusted corresponding tests to validate the new behavior in both classes.

laborious/activities/gates.py

@@ -656,28 +656,29 @@ class Gates(SientiaMonitoring):
 
         for server_id, tags in opc_metrics.items():
             for tag, response_time in tags.items():
-                await self.emit_metric(
+                if response_time is not None:
-                    metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
+                    await self.emit_metric(
-                    method='observe',
+                        metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
-                    tags={
+                        method='observe',
-                        'pod_id': self.pod_id,
+                        tags={
-                        'model_name': metadata['model_name'],
+                            'pod_id': self.pod_id,
-                        'workflow_name': metadata['workflow_name'],
+                            'model_name': metadata['model_name'],
-                        'opc_server_id': server_id,
+                            'workflow_name': metadata['workflow_name'],
-                        'tag': tag,
+                            'opc_server_id': server_id,
-                    },
+                            'tag': tag,
-                    value=response_time,
+                        },
-                )
+                        value=response_time,
+                    )
 
-                await self.emit_metric(
+                    await self.emit_metric(
-                    metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
+                        metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
-                    tags={
+                        tags={
-                        'pod_id': self.pod_id,
+                            'pod_id': self.pod_id,
-                        'model_name': metadata['model_name'],
+                            'model_name': metadata['model_name'],
-                        'workflow_name': metadata['workflow_name'],
+                            'workflow_name': metadata['workflow_name'],
-                        'opc_server_id': server_id,
+                            'opc_server_id': server_id,
-                        'tag': tag,
+                            'tag': tag,
-                    },
+                        },
-                )
+                    )
 
         self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

laborious/activities/mlflow.py

@@ -312,9 +312,12 @@ class MLFlow(SientiaMonitoring):
         self.debug(f'Timestamp: {timestamp}', metadata)
 
         # Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
-        data = data.sort_values('created_at', ascending=False).drop_duplicates(
+        if 'created_at' in data.columns:
-            subset=['variable', 'timestamp'], keep='first'
+            data = data.sort_values('created_at', ascending=False).drop_duplicates(
-        )
+                subset=['variable', 'timestamp'], keep='first'
+            )
+        else:
+            data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
 
         data.drop(columns=['model_id'], inplace=True, errors='ignore')
         data.drop(columns=['created_at'], inplace=True, errors='ignore')

tests/laborious/activities/test_mlflow.py

@@ -254,11 +254,11 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
 
     timestamp = raw_data.__getitem__.return_value.max.return_value
 
-    raw_data.sort_values.assert_called_once_with('created_at', ascending=False)
+    raw_data.sort_values.assert_not_called()
-    raw_data.sort_values.return_value.drop_duplicates.assert_called_once_with(
+    raw_data.drop_duplicates.assert_called_once_with(
         subset=['variable', 'timestamp'], keep='first'
     )
-    raw_data = raw_data.sort_values.return_value.drop_duplicates.return_value
+    raw_data = raw_data.drop_duplicates.return_value
 
     raw_data.drop.assert_has_calls(
         [
@@ -313,7 +313,9 @@ async def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow)
         'message': 'Model retrained failed.',
     }
 
-    mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock()
+    mlflow.minio_repository.get_parquet_as_dataframe.return_value = MagicMock(
+        columns=['variable', 'timestamp', 'value', 'created_at']
+    )
 
     response = await mlflow.retrain_model(
         {

tests/laborious/utils/repository/test_minio_repository.py

@@ -88,6 +88,21 @@ async def test_create_bucket_success(minio_repository):
     )
 
 
+@mark.asyncio
+async def test_create_bucket_error(minio_repository):
+    minio_repository.s3_client.create_bucket.side_effect = ValueError('test')
+
+    with raises(ValueError):
+        await minio_repository.create_bucket({})
+
+
+        minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
+        minio_repository.emit_metric.assert_called_once_with(
+            metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY
+        )
+        minio_repository.observe_lag.assert_not_called()
+
+
 @mark.asyncio
 async def test_ensure_bucket_exists_bucket_exists(minio_repository):
     assert await minio_repository.ensure_bucket_exists({}) is None

tests/laborious/utils/repository/test_model_repository.py

@@ -286,6 +286,16 @@ def test_get_experiment_error(mlflow, mlflow_repository):
         raise AssertionError('Expected ValueError')
 
 
+def test_get_experiment_create_error(mlflow, mlflow_repository):
+    mlflow.get_experiment_by_name.return_value = None
+    mlflow.create_experiment.return_value = None
+    mlflow.get_experiment.return_value = None
+    with pytest.raises(ValueError) as e:
+        mlflow_repository.get_experiment('test', create_if_not_exists=True)
+
+        assert str(e) == 'Experiment test not found after creation, unknown reason'
+
+
 @pytest.mark.asyncio
 async def test_load_predict_model_sklearn(mlflow, mlflow_repository):
     result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn')