SIENTIAPDE-1717: Remove MinIO cleanup functionality and associated components. This change streamlines the cleanup workflow to focus solely on local temporary directories, removes the ModelTrainingError exception, and updates related configurations, documentation, and tests.

This commit is contained in:
Bruno Domingues
2026-03-30 14:14:10 -03:00
parent 63a94dae0a
commit 7a0961f29d
20 changed files with 56 additions and 778 deletions

View File

@@ -8,8 +8,8 @@ import pytest
@pytest.mark.asyncio
@patch('model_manager.workflows.cleanup_files.workflow')
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
"""Test the CleanupFiles workflow when bucket_name is provided in the input."""
async def test_cleanup_files_workflow(mock_workflow_module):
"""Test the CleanupFiles workflow."""
from model_manager.workflows.cleanup_files import CleanupFiles
# Mock execute_activity_method
@@ -17,58 +17,14 @@ async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
# Instantiate and run the workflow
workflow_instance = CleanupFiles()
await workflow_instance.run({'bucket_name': 'input-bucket'})
await workflow_instance.run({})
# Verify that the activities were called with the correct parameters
calls = mock_workflow_module.execute_activity_method.call_args_list
assert len(calls) == 2
# Check cleanup_minio_files call
minio_call_args = calls[0][0][1]
assert minio_call_args['bucket_name'] == 'input-bucket'
assert minio_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
assert len(calls) == 1
# Check cleanup_temp_directories call
local_call_args = calls[1][0][1]
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
assert local_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
@pytest.mark.asyncio
@patch('model_manager.workflows.cleanup_files.workflow')
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
@patch('model_manager.workflows.cleanup_files.DEFAULT_CLEANUP_BUCKET', 'env-var-bucket')
async def test_cleanup_files_workflow_with_default_bucket(mock_workflow_module):
"""Test the CleanupFiles workflow when using the default bucket from environment variables."""
from model_manager.workflows.cleanup_files import CleanupFiles
# Mock execute_activity_method
mock_workflow_module.execute_activity_method = AsyncMock()
# Instantiate and run the workflow
workflow_instance = CleanupFiles()
await workflow_instance.run({}) # Empty input
# Verify that the activities were called
calls = mock_workflow_module.execute_activity_method.call_args_list
assert len(calls) == 2
# Check cleanup_minio_files call
minio_call_args = calls[0][0][1]
assert minio_call_args['bucket_name'] == 'env-var-bucket'
assert minio_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
# Check cleanup_temp_directories call
local_call_args = calls[1][0][1]
local_call_args = calls[0][0][1]
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
assert local_call_args['metadata'] == {
'pod_id': 'temporal-pod',

View File

@@ -4,7 +4,6 @@ from unittest.mock import AsyncMock, Mock, patch
import pytest
from model_manager.utils.exceptions import ModelTrainingError
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
@@ -275,20 +274,18 @@ async def test_train_model_mlflow_error(mock_workflow_module, mock_train_params)
from model_manager.workflows.train_model import TrainModel
# Setup mocks - MLflow save fails
mlflow_error = ModelTrainingError(
model_trained=True, model_saved=False, message='MLflow save failed'
)
mlflow_error = RuntimeError('MLflow save failed')
mock_workflow_module.execute_activity_method = AsyncMock(side_effect=[mlflow_error, None])
workflow_instance = TrainModel()
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
with pytest.raises(ModelTrainingError):
with pytest.raises(RuntimeError):
await workflow_instance._train_model(mock_train_params, 123, metadata)
# Verify TRACKING_SEND_ERROR status was set
# Verify TRAINING_ERROR status was set
call_args = mock_workflow_module.execute_activity_method.call_args_list[1]
assert call_args[0][1]['status'] == ExperimentStatus.TRACKING_SEND_ERROR
assert call_args[0][1]['status'] == ExperimentStatus.TRAINING_ERROR
@pytest.mark.asyncio
@@ -306,14 +303,11 @@ async def test_cleanup_resources_success(mock_workflow_module):
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
await workflow_instance._cleanup_resources(
experiment_run_id=123,
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
)
assert mock_workflow_module.execute_activity_method.call_count == 2
assert mock_workflow_module.execute_activity_method.call_count == 1
@pytest.mark.asyncio
@@ -332,15 +326,12 @@ async def test_cleanup_resources_failure(mock_workflow_module):
with pytest.raises(RuntimeError, match='Cleanup failed'):
await workflow_instance._cleanup_resources(
experiment_run_id=123,
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
)
# Verify error status update was called
assert mock_workflow_module.execute_activity_method.call_count == 2
assert mock_workflow_module.execute_activity_method.call_count == 1
@pytest.mark.asyncio
@@ -411,7 +402,7 @@ async def test_update_experiment_run_with_run_name(mock_workflow_module):
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.TRACKING_SENT,
status=ExperimentStatus.TRAINING_SUCCESS,
run_name='test-run-123',
)
@@ -440,9 +431,8 @@ async def test_run_complete_workflow_success(
mock_train_params, # validate_train_params
None, # update status (ORCHESTRATOR_WAITING_PROC)
train_result, # train_model
None, # update status (TRACKING_SENT)
None, # update status (TRAINING_SUCCESS)
None, # cleanup_resources
None, # update status (FILE_DELETED)
]
)
@@ -452,7 +442,7 @@ async def test_run_complete_workflow_success(
await workflow_instance.run(sample_input_data)
# Verify all activities were called
assert mock_workflow_module.execute_activity_method.call_count == 6
assert mock_workflow_module.execute_activity_method.call_count == 5
@pytest.mark.asyncio
@@ -523,7 +513,6 @@ async def test_run_workflow_cleanup_error(
mock_train_params, # validate_train_params
None, # update status (ORCHESTRATOR_WAITING_PROC)
train_result, # train_model
None, # update status (TRACKING_SENT)
RuntimeError('Cleanup failed'), # cleanup_resources fails
None, # update status (FILE_DELETE_ERROR)
]
@@ -534,7 +523,7 @@ async def test_run_workflow_cleanup_error(
with pytest.raises(RuntimeError, match='Cleanup failed'):
await workflow_instance.run(sample_input_data)
assert mock_workflow_module.execute_activity_method.call_count == 6
assert mock_workflow_module.execute_activity_method.call_count == 5
@pytest.mark.asyncio