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.

tests/activities/test_training.py

@@ -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

tests/schedules/test_cleanup_schedule.py

@@ -125,6 +125,29 @@ async def test_schedule_exists_handles_exception(mock_temporal_client, mock_logg
     assert 'Error checking if schedule exists' in mock_logger.custom_error.call_args[0][0]
 
 
+@pytest.mark.asyncio
+async def test_needs_schedule_reconcile_handles_describe_exception(mock_logger, metadata):
+    """Test _needs_schedule_reconcile returns True and logs when describe fails."""
+    from model_manager.schedules.cleanup_schedule import _needs_schedule_reconcile
+
+    handle = AsyncMock()
+    handle.describe = AsyncMock(side_effect=RuntimeError('describe failed'))
+
+    needs_reconcile = await _needs_schedule_reconcile(
+        schedule_handle=handle,
+        cleanup_task_queue='cleanup_files-model-manager-worker-queue',
+        logger=mock_logger,
+        metadata=metadata,
+    )
+
+    assert needs_reconcile is True
+    mock_logger.custom_error.assert_called_once()
+    assert (
+        'Error describing cleanup schedule for reconcile'
+        in mock_logger.custom_error.call_args[0][0]
+    )
+
+
 # --- create_cleanup_schedule Tests ---
 
 

tests/sientia/test_exceptions.py

@@ -0,0 +1,9 @@
+"""Unit tests for custom exception aliases."""
+
+from mlflow.exceptions import MlflowException
+
+from model_manager.sientia.exceptions import SientiaMlException
+
+
+def test_sientia_ml_exception_is_mlflow_exception_alias():
+    assert SientiaMlException is MlflowException

tests/utils/repository/test_data_manager_repository.py

@@ -27,9 +27,11 @@ def test_train_test_split_dataframe_no_shuffle():
     assert list(tr['a']) == [0, 1, 2, 3, 4]
 
 
-def test_train_test_split_ndarray():
-    arr = np.arange(20).reshape(10, 2)
-    tr, te = dmr.train_test_split(arr, train_size=0.5, shuffle=False, random_state=None)
+def test_train_test_split_dataframe_returns_dataframes():
+    df = pd.DataFrame(np.arange(20).reshape(10, 2), columns=['a', 'b'])
+    tr, te = dmr.train_test_split(df, train_size=0.5, shuffle=False, random_state=None)
+    assert isinstance(tr, pd.DataFrame)
+    assert isinstance(te, pd.DataFrame)
     assert tr.shape[0] == 5 and te.shape[0] == 5
 
 

tests/worker/test_prepare_worker.py

@@ -3,6 +3,12 @@
 from unittest.mock import MagicMock, patch
 
 
+def test_build_queue_name_without_runtime_uses_default_suffix():
+    from model_manager.worker.prepare_worker import build_queue_name
+
+    assert build_queue_name('TrainModel') == 'train_model-queue'
+
+
 def test_prepare_worker_train_queue_uses_train_limits():
     from model_manager.worker.prepare_worker import prepare_worker
     from model_manager.workflows.train_model import TrainModel

tests/workflows/test_train_model.py

@@ -290,9 +290,11 @@ async def test_run_training_failure_skips_cleanup_activity(
 
 
 @pytest.mark.asyncio
-async def test_run_missing_experiment_run_id():
+@patch('model_manager.workflows.train_model.workflow')
+async def test_run_missing_experiment_run_id(mock_wf):
     from model_manager.workflows.train_model import TrainModel
 
+    mock_wf.logger = Mock()
     with pytest.raises(ValueError, match='experiment_run_id is required'):
         await TrainModel().run({})