SIENTIAPDE-1273

Enhance MLFlowRepository and Activities classes with new methods and metrics

- Added `check_artifact_exists` method to MLFlowRepository for verifying artifact presence in the MLflow Model Registry.
- Implemented `get_prediction_data` method in MLFlowRepository to retrieve prediction data from models.
- Updated Activities class to integrate ModelMetrics for improved metrics handling.
- Enhanced tests for artifact existence checks and prediction data retrieval, ensuring robust coverage for new functionalities.
- Updated various workflows to include `transform_table_name` in input data for better data handling.

tests/laborious/utils/repository/test_model_repository.py

@@ -190,6 +190,26 @@ def test_get_model_params(mlflow, mlflow_repository):
     assert output == mlflow.get_run.return_value.data.params
 
 
+def test_check_artifact_exists_true(mlflow_repository):
+    artifact = MagicMock(path='test_artifact')
+    mlflow_repository.client.list_artifacts.return_value = [artifact]
+    
+    result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
+    
+    assert result is True
+    mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
+
+
+def test_check_artifact_exists_false(mlflow_repository):
+    artifact = MagicMock(path='other_artifact')
+    mlflow_repository.client.list_artifacts.return_value = [artifact]
+    
+    result = mlflow_repository.check_artifact_exists('run_id', 'test_artifact', metadata['metadata'])
+    
+    assert result is False
+    mlflow_repository.client.list_artifacts.assert_called_once_with('run_id')
+
+
 @pytest.mark.asyncio
 @patch('laborious.utils.repository.model_repository.path')
 @patch('laborious.utils.repository.model_repository.rmtree')
@@ -275,6 +295,56 @@ async def test_download_artifacts_error(makedirs, rmtree, path, mlflow_repositor
         mlflow_repository.observe_lag.assert_not_called()
 
 
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.model_repository.mlflow')
+@patch('laborious.utils.repository.model_repository.pd')
+@patch('laborious.utils.repository.model_repository.StringIO')
+async def test_load_artifact_dataframe_success(StringIO, pd, mlflow, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
+    
+    mlflow.artifacts.load_text.return_value = 'col1,col2\n1,2\n3,4'
+    
+    result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+    
+    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
+    mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
+    mlflow.artifacts.load_text.assert_called_once_with('runs:/run_id/artifact_path')
+    assert result == pd.read_csv.return_value
+    mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_READ_LAG, ANY)
+    mlflow_repository.emit_metric.assert_called_once_with(
+        metric_object=metrics.MODEL_READ_COUNT, tags=ANY
+    )
+
+
+@pytest.mark.asyncio
+async def test_load_artifact_dataframe_not_exists(mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=False)
+    
+    result = await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+    
+    assert result is None
+    mlflow_repository.get_model_run_id.assert_called_once_with(model_name='model_name', stage='Production')
+    mlflow_repository.check_artifact_exists.assert_called_once_with('run_id', 'artifact_path', metadata['metadata'])
+
+
+@pytest.mark.asyncio
+@patch('laborious.utils.repository.model_repository.mlflow')
+async def test_load_artifact_dataframe_error(mlflow, mlflow_repository):
+    mlflow_repository.get_model_run_id = MagicMock(return_value='run_id')
+    mlflow_repository.check_artifact_exists = MagicMock(return_value=True)
+    mlflow.artifacts.load_text.side_effect = ValueError('error')
+    
+    with pytest.raises(ValueError):
+        await mlflow_repository.load_artifact_dataframe('model_name', 'artifact_path', metadata['metadata'])
+    
+    mlflow_repository.emit_metric.assert_called_once_with(
+        metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=ANY
+    )
+    mlflow_repository.observe_lag.assert_not_called()
+
+
 def test_get_experiment_error(mlflow, mlflow_repository):
     mlflow.get_experiment_by_name.return_value = None
 
@@ -719,6 +789,7 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
     mlflow_repository.detect_and_parse_datetime_index = MagicMock(
         return_value=MagicMock(drop_duplicates=MagicMock(return_value=MagicMock(columns=[])))
     )
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
 
     data = MagicMock()
 
@@ -775,9 +846,14 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
 
     prediction_model.fit.assert_called_once_with(pd_merge.return_value)
 
+    mlflow_repository.get_prediction_data.assert_called_once_with(
+        prediction_model, pd_merge.return_value, data_model.fit.return_value.target_variable
+    )
+
     assert output == {
         'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'},
         'data_model': {'model': data_model.fit.return_value, 'artifact_path': 'artifact_path'},
+        'prediction_data': mlflow_repository.get_prediction_data.return_value,
     }
 
 
@@ -801,6 +877,7 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
             drop_duplicates=MagicMock(return_value=MagicMock(columns=['feat_1']))
         )
     )
+    mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
 
     data = MagicMock()
 
@@ -855,9 +932,14 @@ async def test_fit_models_df_target_name_not_none_and_in_model(
 
     prediction_model.fit.assert_called_once_with(transformed_data)
 
+    mlflow_repository.get_prediction_data.assert_called_once_with(
+        prediction_model, transformed_data, 'feat_1'
+    )
+
     assert output == {
         'prediction_model': {'model': prediction_model, 'artifact_path': 'artifact_path'},
         'data_model': {'model': data_model, 'artifact_path': 'artifact_path'},
+        'prediction_data': mlflow_repository.get_prediction_data.return_value,
     }
 
 
@@ -878,7 +960,8 @@ async def test_log_model_sklearn(mlflow, mlflow_repository):
 @patch('laborious.utils.repository.model_repository.path')
 @pytest.mark.asyncio
 async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
-    model_data = {'model': MagicMock(), 'artifact_path': 'artifact_path'}
+    model_mock = MagicMock()
+    model_data = {'model': model_mock, 'artifact_path': 'artifact_path'}
     await mlflow_repository.log_model(
         model_data, 'pyfunc', 'prediction_model', metadata['metadata']
     )
@@ -887,7 +970,7 @@ async def test_log_model_pyfunc(path, mlflow, mlflow_repository):
 
     path.join.assert_called_once_with('artifact_path', 'code', 'utils')
 
-    model_data['model'].store_model.assert_called_once_with(
+    model_mock.store_model.assert_called_once_with(
         artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False
     )
 
@@ -934,9 +1017,11 @@ async def test_create_new_experiment(
 ):
     model_name = 'model_name'
     data = MagicMock()
+    prediction_data = MagicMock(spec=DataFrame)
     retrain_data = {
         'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
         'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
+        'prediction_data': prediction_data,
     }
 
     mlflow_repository.get_model_params = MagicMock(
@@ -971,7 +1056,11 @@ async def test_create_new_experiment(
         mlflow_repository.get_experiment.return_value.name
     )
 
-    data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=True)
+    data.to_csv.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv', index=False)
+    # Verify prediction_data.to_csv was called with correct arguments
+    prediction_data.to_csv.assert_called_once()
+    assert prediction_data.to_csv.call_args[0][0] == './tmp/artifacts/model_name/evaluation_data.csv'
+    assert prediction_data.to_csv.call_args[1]['index'] is False
 
     mlflow.start_run.assert_called_once_with(
         experiment_id=mlflow_repository.get_experiment.return_value.experiment_id,
@@ -1000,7 +1089,10 @@ async def test_create_new_experiment(
         }
     )
 
-    mlflow.log_artifact.assert_called_once_with('./tmp/artifacts/model_name/retrain_data.csv')
+    mlflow.log_artifact.assert_has_calls([
+        call('./tmp/artifacts/model_name/retrain_data.csv'),
+        call('./tmp/artifacts/model_name/evaluation_data.csv'),
+    ])
 
     force_memory_release.assert_called_once_with(mlflow_repository.logger)
 
@@ -1026,9 +1118,11 @@ async def test_create_new_experiment_error(
     mlflow.start_run.side_effect = ValueError('error')
     model_name = 'model_name'
     data = MagicMock()
+    prediction_data = MagicMock(spec=DataFrame)
     retrain_data = {
         'prediction_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
         'data_model': {'model': MagicMock(), 'artifact_path': 'artifact_path'},
+        'prediction_data': prediction_data,
     }
 
     mlflow_repository.get_model_params = MagicMock(
@@ -1371,3 +1465,33 @@ async def test_update_production_model(mlflow_repository):
         'mlflow_run_id': '0',
         'mlflow_experiment_id': '0',
     }
+
+
+def test_get_prediction_data_dataframe(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = DataFrame({'pred': [5, 6]}, index=['idx1', 'idx2'])
+    target_name = 'target'
+    
+    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    
+    prediction_model.predict.assert_called_once_with(retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]
+
+
+def test_get_prediction_data_array(mlflow_repository):
+    prediction_model = MagicMock()
+    retrain_dataset = DataFrame({'feat_1': [1, 2], 'target': [3, 4]}, index=['idx1', 'idx2'])
+    prediction_model.predict.return_value = [5, 6]
+    target_name = 'target'
+    
+    result = mlflow_repository.get_prediction_data(prediction_model, retrain_dataset, target_name)
+    
+    prediction_model.predict.assert_called_once_with(retrain_dataset)
+    assert 'prediction' in result.columns
+    assert 'target' in result.columns
+    assert 'timestamp' in result.columns
+    assert result.index.tolist() == [0, 1]