feat: update training workflow and repository management
- Replaced synchronous MinIO repository calls with asynchronous counterparts in the Training class for improved performance. - Enhanced logging throughout the training process to provide better insights into model metadata loading, parameter validation, and training execution. - Updated the train_test_split function to enforce DataFrame input type, ensuring consistency in data handling. - Removed the deprecated model_repository.py file to streamline the codebase. - Adjusted cleanup schedule logic to improve error handling and logging during schedule reconciliation. - Updated tests to reflect changes in the training workflow and repository interactions.
This commit is contained in:
@@ -93,7 +93,7 @@ def test_train_model_download_fails_notifies(training):
|
||||
'model_metadata': {'schemas': {'components': {'schemas': {}}}},
|
||||
}
|
||||
)
|
||||
training.minio_repository.download_file_sync = MagicMock(side_effect=OSError('minio'))
|
||||
training.minio_repository.download_file = MagicMock(side_effect=OSError('minio'))
|
||||
training.send_notification = MagicMock()
|
||||
with pytest.raises(OSError, match='minio'):
|
||||
training.train_model({'metadata': {'pod': 'x'}, 'train_params': tp.to_dict()})
|
||||
@@ -129,7 +129,7 @@ def test_train_model_success_serializes_result(mock_mlflow, training):
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'csv')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: setattr(x, 'mse_val', 0.1) or x
|
||||
@@ -185,7 +185,7 @@ def test_train_model_train_params_as_dict(mock_mlflow, training):
|
||||
tp = TrainModelParams.from_dict(d)
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'csv')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'csv')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
@@ -240,7 +240,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
||||
return b'val'
|
||||
raise AssertionError(object_name)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(side_effect=_dl)
|
||||
training.minio_repository.download_file = MagicMock(side_effect=_dl)
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
@@ -272,7 +272,7 @@ def test_train_model_downloads_validation_file_when_set(mock_mlflow, training):
|
||||
training.mlflow_repository.start_run = _run_ctx
|
||||
|
||||
training.train_model({'metadata': {}, 'train_params': tp.to_dict()})
|
||||
assert training.minio_repository.download_file_sync.call_count == 2
|
||||
assert training.minio_repository.download_file.call_count == 2
|
||||
mock_mlflow.log_artifact.assert_called()
|
||||
|
||||
|
||||
@@ -288,7 +288,7 @@ def test_train_model_value_error_when_paths_missing_after_report(training):
|
||||
val_df = pd.DataFrame({'a': [1.0], 't': [1.0]})
|
||||
tmr = TrainModelResult(params=tp, train_data=train_df, val_data=val_df)
|
||||
|
||||
training.minio_repository.download_file_sync = MagicMock(return_value=b'x')
|
||||
training.minio_repository.download_file = MagicMock(return_value=b'x')
|
||||
training.data_manager_repository.prepare_training_data = MagicMock(return_value=tmr)
|
||||
training.data_manager_repository.compute_regression_metrics = MagicMock(
|
||||
side_effect=lambda x, _w: x
|
||||
|
||||
Reference in New Issue
Block a user