diff --git a/model_manager/activities/mlflow.py b/model_manager/activities/mlflow.py index 1762b04..b01b48a 100644 --- a/model_manager/activities/mlflow.py +++ b/model_manager/activities/mlflow.py @@ -429,3 +429,67 @@ class MLFlow(BaseActivity): # Re-raise exception to stop workflow raise + + @activity.defn(name='cleanup_run_directory') + async def cleanup_run_directory(self, input_data: dict[str, Any]) -> None: + """ + Clean up temporary run directory after model training. + + This activity deletes the temporary directory created during model training + and artifact generation. It implements idempotent cleanup to handle cases + where the directory may have already been deleted. + + Args: + input_data: Configuration for cleanup operation + Required keys: + - metadata (dict): Workflow execution metadata + - run_dir (str): Path to the run directory to delete + + Raises: + Exception: If cleanup fails for reasons other than directory not existing + + Example: + await cleanup_run_directory({ + 'metadata': {'workflow_id': 'cleanup-123'}, + 'run_dir': '/path/to/run_dir' + }) + """ + import os + import shutil + + metadata = input_data.get('metadata', {}) + run_dir = input_data.get('run_dir') + + try: + if not run_dir: + self.info('No run directory specified, skipping cleanup', metadata) + return + + self.info(f'Cleaning up run directory: {run_dir}', metadata) + + # Idempotent cleanup: check if directory exists before deleting + if os.path.exists(run_dir): + shutil.rmtree(run_dir) + self.info(f'Run directory deleted successfully: {run_dir}', metadata) + else: + self.info(f'Run directory already deleted: {run_dir}', metadata) + + except Exception as e: + error_msg = f'Error cleaning up run directory {run_dir}: {str(e)}' + trace = traceback.format_exc() + + # Send notification + self.send_notification( + metadata=metadata, + notification_id='CLEANUP_RUN_DIRECTORY_ERROR', + message=error_msg, + block='cleanup_run_directory', + level=NotificationLevel.ERROR, + attachment_content=trace, + ) + + # Log error + self.error(trace, metadata=metadata) + + # Re-raise exception + raise diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index c3a4065..1d55884 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -12,7 +12,6 @@ This workflow orchestrates the complete ML model training process, including: from temporalio import workflow with workflow.unsafe.imports_passed_through(): - import shutil from datetime import timedelta from typing import Any @@ -373,8 +372,8 @@ class TrainModel: """ Cleanup resources and delete file from MinIO. - This method removes the temporary run directory and deletes the training - file from MinIO. On success, updates DB status to FILE_DELETED. + This method removes the temporary run directory via activity and deletes + the training file from MinIO. On success, updates DB status to FILE_DELETED. On error, updates DB status to FILE_DELETE_ERROR. Args: @@ -386,11 +385,22 @@ class TrainModel: Exception: If cleanup fails (after updating DB status) """ try: - # Step 1: Remove temporary run directory + # Step 1: Remove temporary run directory via activity (deterministic) if hasattr(saved_result, 'run_dir') and saved_result.run_dir: - shutil.rmtree(saved_result.run_dir) + cleanup_input = { + **metadata, + 'run_dir': saved_result.run_dir, + } + + await workflow.execute_activity_method( + Activities.cleanup_run_directory, + cleanup_input, + retry_policy=retry_policy, + start_to_close_timeout=timedelta(seconds=30), + ) + workflow.logger.info( - f'Removed run directory: {saved_result.run_dir} for experiment {experiment_run_id}' + f'Run directory cleanup completed for experiment {experiment_run_id}' ) # Step 2: Delete file from MinIO diff --git a/tests/activities/test_mlflow.py b/tests/activities/test_mlflow.py index 7a0c1cb..0c965ec 100644 --- a/tests/activities/test_mlflow.py +++ b/tests/activities/test_mlflow.py @@ -532,3 +532,165 @@ async def test_save_model_complete_flow(mlflow): assert response.run_name == 'production_model-5' assert response.run_dir is not None assert response.report_path is not None + + +# ============================================================================ +# Tests for cleanup_run_directory +# ============================================================================ + + +@mark.asyncio +async def test_cleanup_run_directory_success(mlflow): + """Test successful cleanup of run directory.""" + input_data = { + **metadata, + 'run_dir': 'test_run_dir', + } + + # Patch os and shutil inside the activity method + with ( + patch('os.path.exists', return_value=True) as mock_exists, + patch('shutil.rmtree') as mock_rmtree, + ): + # Call the method + await mlflow.cleanup_run_directory(input_data) + + # Verify directory existence was checked + mock_exists.assert_called_once_with('test_run_dir') + + # Verify shutil.rmtree was called + mock_rmtree.assert_called_once_with('test_run_dir') + + +@mark.asyncio +async def test_cleanup_run_directory_already_deleted(mlflow): + """Test cleanup when directory is already deleted (idempotent).""" + input_data = { + **metadata, + 'run_dir': 'already_deleted_dir', + } + + # Patch os and shutil inside the activity method + with ( + patch('os.path.exists', return_value=False) as mock_exists, + patch('shutil.rmtree') as mock_rmtree, + ): + # Call the method - should not raise error + await mlflow.cleanup_run_directory(input_data) + + # Verify directory existence was checked + mock_exists.assert_called_once_with('already_deleted_dir') + + # Verify shutil.rmtree was NOT called + mock_rmtree.assert_not_called() + + +@mark.asyncio +async def test_cleanup_run_directory_no_run_dir(mlflow): + """Test cleanup when no run_dir is provided.""" + input_data = { + **metadata, + # No run_dir key + } + + # Call the method - should not raise error + await mlflow.cleanup_run_directory(input_data) + + # Should complete without errors + + +@mark.asyncio +async def test_cleanup_run_directory_empty_run_dir(mlflow): + """Test cleanup when run_dir is empty string.""" + input_data = { + **metadata, + 'run_dir': '', + } + + # Call the method - should not raise error + await mlflow.cleanup_run_directory(input_data) + + # Should complete without errors + + +@mark.asyncio +async def test_cleanup_run_directory_none_run_dir(mlflow): + """Test cleanup when run_dir is None.""" + input_data = { + **metadata, + 'run_dir': None, + } + + # Call the method - should not raise error + await mlflow.cleanup_run_directory(input_data) + + # Should complete without errors + + +@mark.asyncio +async def test_cleanup_run_directory_error(mlflow): + """Test cleanup handles errors correctly.""" + input_data = { + **metadata, + 'run_dir': 'error_dir', + } + + # Patch os and shutil inside the activity method + with ( + patch('os.path.exists', return_value=True), + patch('shutil.rmtree', side_effect=PermissionError('Permission denied')), + ): + # Call the method - should raise exception + with pytest.raises(PermissionError, match='Permission denied'): + await mlflow.cleanup_run_directory(input_data) + + # Verify notification was sent + mlflow.send_notification.assert_called_once_with( + metadata=metadata['metadata'], + notification_id='CLEANUP_RUN_DIRECTORY_ERROR', + message=ANY, + block='cleanup_run_directory', + level=NotificationLevel.ERROR, + attachment_content=ANY, + ) + + +@mark.asyncio +async def test_cleanup_run_directory_missing_metadata(mlflow): + """Test cleanup handles missing metadata gracefully.""" + input_data = { + 'run_dir': 'no_metadata_dir', + } + + # Patch os and shutil inside the activity method + with patch('os.path.exists', return_value=True), patch('shutil.rmtree') as mock_rmtree: + # Call the method + await mlflow.cleanup_run_directory(input_data) + + # Verify directory was deleted + mock_rmtree.assert_called_once_with('no_metadata_dir') + + +@mark.asyncio +async def test_cleanup_run_directory_oserror(mlflow): + """Test cleanup handles OSError correctly.""" + input_data = { + **metadata, + 'run_dir': 'os_error_dir', + } + + # Patch os and shutil inside the activity method + with ( + patch('os.path.exists', return_value=True), + patch('shutil.rmtree', side_effect=OSError('Directory not empty')), + ): + # Call the method - should raise exception + with pytest.raises(OSError, match='Directory not empty'): + await mlflow.cleanup_run_directory(input_data) + + # Verify notification was sent + mlflow.send_notification.assert_called_once() + call_args = mlflow.send_notification.call_args[1] + assert call_args['notification_id'] == 'CLEANUP_RUN_DIRECTORY_ERROR' + assert call_args['level'] == NotificationLevel.ERROR + assert 'Directory not empty' in call_args['message'] diff --git a/tests/workflows/test_train_model.py b/tests/workflows/test_train_model.py index 8db760a..561e863 100644 --- a/tests/workflows/test_train_model.py +++ b/tests/workflows/test_train_model.py @@ -59,11 +59,9 @@ def mock_train_result(mock_train_params): @mark.asyncio -@patch('model_manager.workflows.train_model.shutil') @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock) async def test_run_success_complete_flow( workflow_mock: AsyncMock, - mock_shutil: MagicMock, train_model_workflow: TrainModel, mock_train_params, mock_train_result, @@ -101,6 +99,7 @@ async def test_run_success_complete_flow( None, # update_experiment_run (TRAINING_SUCCESS) mock_train_result, # save_model None, # update_experiment_run (MLFLOW_SENT with run_name) + None, # cleanup_run_directory None, # delete_file_from_minio None, # update_experiment_run (FILE_DELETED) ] @@ -109,10 +108,8 @@ async def test_run_success_complete_flow( # Execute workflow await train_model_workflow.run(input_data) - # Verify all activity calls - assert workflow_mock.execute_activity_method.call_count == 9 - # Verify shutil.rmtree was called - mock_shutil.rmtree.assert_called_once_with('test_run_dir') + # Verify all activity calls (now 10 instead of 9 due to cleanup_run_directory) + assert workflow_mock.execute_activity_method.call_count == 10 @mark.asyncio @@ -506,9 +503,7 @@ async def test_save_model_to_mlflow_error( @mark.asyncio @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock) -@patch('model_manager.workflows.train_model.shutil') async def test_cleanup_resources_success( - mock_shutil: MagicMock, workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result, @@ -518,6 +513,7 @@ async def test_cleanup_resources_success( workflow_mock.execute_activity_method = AsyncMock( side_effect=[ + None, # cleanup_run_directory None, # delete_file_from_minio None, # update_experiment_run ] @@ -527,18 +523,13 @@ async def test_cleanup_resources_success( saved_result=mock_train_result, experiment_run_id=123, metadata=metadata ) - # Verify shutil.rmtree was called - mock_shutil.rmtree.assert_called_once_with('test_run_dir') - - # Verify activities were called - assert workflow_mock.execute_activity_method.call_count == 2 + # Verify activities were called (cleanup_run_directory + delete_file_from_minio + update_experiment_run) + assert workflow_mock.execute_activity_method.call_count == 3 @mark.asyncio @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock) -@patch('model_manager.workflows.train_model.shutil') async def test_cleanup_resources_delete_error( - mock_shutil: MagicMock, workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result, @@ -548,6 +539,7 @@ async def test_cleanup_resources_delete_error( workflow_mock.execute_activity_method = AsyncMock( side_effect=[ + None, # cleanup_run_directory succeeds Exception('MinIO delete failed'), # delete_file_from_minio fails None, # update_experiment_run with error ] @@ -558,15 +550,13 @@ async def test_cleanup_resources_delete_error( saved_result=mock_train_result, experiment_run_id=123, metadata=metadata ) - # Verify error status was updated - assert workflow_mock.execute_activity_method.call_count == 2 + # Verify error status was updated (cleanup_run_directory + delete_file_from_minio + update_experiment_run) + assert workflow_mock.execute_activity_method.call_count == 3 @mark.asyncio @patch('model_manager.workflows.train_model.workflow', new_callable=AsyncMock) -@patch('model_manager.workflows.train_model.shutil') async def test_cleanup_resources_without_run_dir( - mock_shutil: MagicMock, workflow_mock: AsyncMock, train_model_workflow: TrainModel, mock_train_result, @@ -588,10 +578,8 @@ async def test_cleanup_resources_without_run_dir( saved_result=mock_train_result, experiment_run_id=123, metadata=metadata ) - # Verify shutil.rmtree was NOT called - mock_shutil.rmtree.assert_not_called() - - # Verify activities were still called + # Verify cleanup_run_directory was NOT called (no run_dir) + # Only delete_file_from_minio + update_experiment_run assert workflow_mock.execute_activity_method.call_count == 2