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,6 +656,7 @@ class Gates(SientiaMonitoring):
 
         for server_id, tags in opc_metrics.items():
             for tag, response_time in tags.items():
+                if response_time is not None:
                     await self.emit_metric(
                         metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
                         method='observe',

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
+        if 'created_at' in data.columns:
             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.return_value.drop_duplicates.assert_called_once_with(
+    raw_data.sort_values.assert_not_called()
+    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')