SIENTIAPDE-1253: Implement activity for cleaning up temporary run directory and integrate into train model workflow. This change introduces a new activity for idempotent cleanup of the run directory after model training, replacing the direct directory removal in the workflow. This improves determinism and error handling. Also includes unit tests for the new activity.
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user