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:
@@ -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',
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user