gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: enhance training and experiment tracking functionality

Browse Source 
- Updated `Activities` class to improve garbage collection handling.
- Enhanced error messaging in `ExperimentTracking` for better clarity on update failures.
- Refactored `Training` class to streamline exception handling and improve type hints.
- Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata.
- Added functionality to extract model equations in `DataManagerRepository` for linear regression models.
This commit is contained in:
vitor-aignosi
2026-04-06 15:05:57 -03:00
parent 1352d1ac8f
commit 6b1df7c3a7
22 changed files with 1751 additions and 2085 deletions
Download Patch File Download Diff File Expand all files Collapse all files

281
tests/worker/test_worker.py
View File

@@ -68,10 +68,11 @@ def mock_activities():
"""Create a mock Activities instance."""
activities = AsyncMock()
activities.update_experiment_run = Mock()
activities.load_model_metadata = Mock()
activities.validate_train_params = Mock()
activities.train_model = Mock()
activities.cleanup_resources = Mock()
activities.shutdown = AsyncMock()
activities.shutdown = Mock()
return activities
@@ -175,7 +176,8 @@ def test_start_prometheus_server_failure(
@pytest.mark.asyncio
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -203,7 +205,7 @@ async def test_main_successful_startup(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_env_vars,
mock_logger,
mock_temporal_client,
@@ -243,19 +245,23 @@ async def test_main_successful_startup(
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(
side_effect=asyncio.CancelledError()
) # Simulate interruption
mock_worker_class.return_value = mock_worker_instance
mock_prepare_worker.return_value = mock_worker_instance
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
@@ -273,7 +279,7 @@ async def test_main_successful_startup(
mock_activities_class.assert_called_once()
mock_client_class.connect.assert_called_once()
# Agora são criados dois Workers: um para train_model-queue e outro para cleanup-queue
assert mock_worker_class.call_count == 2
assert mock_prepare_worker.call_count == 2
# Verify cleanup was performed
mock_notification_handler.shutdown.assert_called_once()
@@ -282,7 +288,8 @@ async def test_main_successful_startup(
@pytest.mark.asyncio
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -310,7 +317,7 @@ async def test_main_handles_exception(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_env_vars,
mock_logger,
):
@@ -333,18 +340,21 @@ async def test_main_handles_exception(
mock_notification_handler_class.return_value = mock_notification_handler
mock_activities = AsyncMock()
mock_activities.shutdown = AsyncMock()
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=RuntimeError('Worker failed'))
mock_worker_class.return_value = mock_worker_instance
mock_prepare_worker.return_value = mock_worker_instance
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
@@ -364,6 +374,7 @@ async def test_main_handles_exception(
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
# Run main() and expect SystemExit
@@ -383,8 +394,9 @@ async def test_main_handles_exception(
@pytest.mark.asyncio
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
@patch('model_manager.worker.worker.create_cleanup_schedule')
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -412,7 +424,7 @@ async def test_main_temporal_client_configuration(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_create_cleanup_schedule,
mock_logger,
):
@@ -459,6 +471,7 @@ async def test_main_temporal_client_configuration(
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
mock_notification_handler = Mock()
@@ -466,18 +479,21 @@ async def test_main_temporal_client_configuration(
mock_notification_handler_class.return_value = mock_notification_handler
mock_activities = AsyncMock()
mock_activities.shutdown = AsyncMock()
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
mock_worker_class.return_value = mock_worker_instance
mock_prepare_worker.return_value = mock_worker_instance
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
@@ -496,8 +512,9 @@ async def test_main_temporal_client_configuration(
@pytest.mark.asyncio
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
@patch('model_manager.worker.worker.create_cleanup_schedule')
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -525,108 +542,104 @@ async def test_main_worker_configuration(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_create_cleanup_schedule,
mock_env_vars,
mock_logger,
):
"""Test that Temporal worker is configured with correct parameters."""
"""Test that prepare_worker is configured with correct workflows and activities."""
from model_manager.worker.worker import main
from model_manager.workflows.cleanup_files import CleanupFiles
from model_manager.workflows.train_model import TrainModel
mock_create_cleanup_schedule.return_value = AsyncMock()
# Patch the task queue constants directly
with (
patch('model_manager.worker.worker.TRAIN_TASK_QUEUE', 'train_model-local_queue'),
patch('model_manager.worker.worker.CLEANUP_TASK_QUEUE', 'cleanup-local_queue'),
):
# Setup mocks
mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
'uri': 'localhost:27018',
}
mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {}
mock_get_logger.return_value = mock_logger
mock_build_mongodb.return_value = {
'connection_string': 'mongodb://test',
'database_name': 'test_db',
'uri': 'localhost:27018',
}
mock_build_postgres.return_value = {}
mock_build_mlflow.return_value = {}
mock_build_minio.return_value = {}
mock_notification_handler = Mock()
mock_notification_handler_class.return_value = mock_notification_handler
mock_notification_handler = Mock()
mock_notification_handler_class.return_value = mock_notification_handler
mock_activities = AsyncMock()
mock_activities.update_experiment_run = Mock()
mock_activities.validate_train_params = Mock()
mock_activities.train_model = Mock()
mock_activities.cleanup_resources = Mock()
mock_activities.shutdown = AsyncMock()
mock_activities_class.return_value = mock_activities
mock_activities = AsyncMock()
mock_activities.update_experiment_run = Mock()
mock_activities.load_model_metadata = Mock()
mock_activities.validate_train_params = Mock()
mock_activities.train_model = Mock()
mock_activities.cleanup_resources = Mock()
mock_activities.cleanup_temp_directories = Mock()
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
mock_worker_class.return_value = mock_worker_instance
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
mock_prepare_worker.return_value = mock_worker_instance
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
mock_plugin_store_instance = AsyncMock()
mock_plugin_store_instance.install_runtime = AsyncMock(
return_value={'runtime': 'model-manager-worker', 'installed': []},
)
mock_plugin_store_class.return_value = mock_plugin_store_instance
mock_build_plugin_store_config.return_value = {
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
'owner': 'sientia',
'repo': 'model-library-store',
'branch': 'main',
'username': 'gitea-user',
'password': 'gitea-password',
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
}
mock_plugin_store_instance = AsyncMock()
mock_plugin_store_instance.install_runtime = AsyncMock(
return_value={'runtime': 'model-manager-worker', 'installed': []},
)
mock_plugin_store_class.return_value = mock_plugin_store_instance
mock_build_plugin_store_config.return_value = {
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
'owner': 'sientia',
'repo': 'model-library-store',
'branch': 'main',
'username': 'gitea-user',
'password': 'gitea-password',
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
# Run main()
with pytest.raises(SystemExit):
await main()
with pytest.raises(SystemExit):
await main()
# Verify Worker was created with correct configuration
assert mock_worker_class.call_count == 2
assert mock_prepare_worker.call_count == 2
# Primeira chamada: worker de treinamento (train_model-local_queue)
train_call_args = mock_worker_class.call_args_list[0]
assert train_call_args[0][0] == mock_client_instance # temporal_client
assert train_call_args[1]['task_queue'] == 'train_model-local_queue'
assert train_call_args[1]['max_concurrent_workflow_tasks'] == 10
assert train_call_args[1]['max_concurrent_activities'] == 10
assert train_call_args[1]['max_concurrent_local_activities'] == 10
assert train_call_args[1]['max_cached_workflows'] == 100
train_call = mock_prepare_worker.call_args_list[0]
assert train_call.kwargs['temporal_client'] is mock_client_instance
assert train_call.kwargs['logger'] is mock_logger
assert train_call.kwargs['main_workflow'] is TrainModel
assert train_call.kwargs['other_workflows'] == []
train_activities_list = train_call.kwargs['activities']
assert mock_activities.update_experiment_run in train_activities_list
assert mock_activities.load_model_metadata in train_activities_list
assert mock_activities.validate_train_params in train_activities_list
assert mock_activities.train_model in train_activities_list
assert mock_activities.cleanup_resources in train_activities_list
train_activities_list = train_call_args[1]['activities']
assert mock_activities.update_experiment_run in train_activities_list
assert mock_activities.validate_train_params in train_activities_list
assert mock_activities.train_model in train_activities_list
assert mock_activities.cleanup_resources in train_activities_list
# Segunda chamada: worker de cleanup (cleanup-local_queue)
cleanup_call_args = mock_worker_class.call_args_list[1]
assert cleanup_call_args[0][0] == mock_client_instance # temporal_client
assert cleanup_call_args[1]['task_queue'] == 'cleanup-local_queue'
assert cleanup_call_args[1]['max_concurrent_workflow_tasks'] == 20
assert cleanup_call_args[1]['max_concurrent_activities'] == 20
assert cleanup_call_args[1]['max_concurrent_local_activities'] == 20
assert cleanup_call_args[1]['max_cached_workflows'] == 100
cleanup_call = mock_prepare_worker.call_args_list[1]
assert cleanup_call.kwargs['temporal_client'] is mock_client_instance
assert cleanup_call.kwargs['logger'] is mock_logger
assert cleanup_call.kwargs['main_workflow'] is CleanupFiles
assert cleanup_call.kwargs['other_workflows'] == []
assert cleanup_call.kwargs['activities'] == [mock_activities.cleanup_temp_directories]
@pytest.mark.asyncio
@patch('model_manager.worker.worker.RUNTIME', 'model-manager-worker')
@patch('model_manager.worker.worker.create_cleanup_schedule')
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -654,7 +667,7 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_create_cleanup_schedule,
mock_logger,
):
@@ -683,18 +696,21 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
mock_notification_handler_class.return_value = mock_notification_handler
mock_activities = AsyncMock()
mock_activities.shutdown = AsyncMock()
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
mock_runtime = Mock()
mock_runtime_class.return_value = mock_runtime
mock_client_instance = AsyncMock()
mock_client_instance.config = Mock(
return_value={'plugins': [], 'interceptors': []}
)
mock_client_class.connect = AsyncMock(return_value=mock_client_instance)
mock_worker_instance = Mock()
mock_worker_instance.run = AsyncMock(side_effect=asyncio.CancelledError())
mock_worker_class.return_value = mock_worker_instance
mock_prepare_worker.return_value = mock_worker_instance
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
@@ -714,11 +730,27 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
'cache_ttl_seconds': None,
}
with pytest.raises(SystemExit):
await main()
mock_create_cleanup_schedule.assert_called_once()
schedule_error_logged = False
for call in mock_logger.custom_error.call_args_list:
if call[0] and 'Failed to configure cleanup schedule' in call[0][0]:
schedule_error_logged = True
break
assert schedule_error_logged, 'Schedule creation error should be logged'
assert mock_prepare_worker.call_count == 2
@pytest.mark.asyncio
@patch('model_manager.worker.worker.Worker')
@patch('model_manager.worker.worker.RUNTIME', None)
@patch('model_manager.worker.worker.prepare_worker')
@patch('model_manager.worker.worker.client.Client')
@patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities')
@@ -746,7 +778,7 @@ async def test_main_missing_runtime_fails_fast(
mock_activities_class,
mock_runtime_class,
mock_client_class,
mock_worker_class,
mock_prepare_worker,
mock_logger,
):
"""Test that main() fails fast when RUNTIME is missing."""
@@ -767,59 +799,14 @@ async def test_main_missing_runtime_fails_fast(
mock_notification_handler_class.return_value = mock_notification_handler
mock_activities = AsyncMock()
mock_activities.shutdown = AsyncMock()
mock_activities.shutdown = Mock()
mock_activities_class.return_value = mock_activities
mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up
mock_plugin_store_instance = AsyncMock()
mock_plugin_store_instance.install_runtime = AsyncMock(
return_value={'runtime': 'model-manager-worker', 'installed': []},
)
mock_plugin_store_class.return_value = mock_plugin_store_instance
mock_build_plugin_store_config.return_value = {
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
'owner': 'sientia',
'repo': 'model-library-store',
'branch': 'main',
'username': 'gitea-user',
'password': 'gitea-password',
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
'pypi_username': None,
'pypi_password': None,
}
# Ensure RUNTIME is not defined
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(SystemExit) as exc_info:
await main()
assert exc_info.value.code == 1
mock_logger.custom_critical.assert_called_once()
mock_app_up.set.assert_called_with(0)
# Run main() - should not fail despite schedule creation error
with pytest.raises(SystemExit):
with pytest.raises(ValueError, match='RUNTIME environment variable is required'):
await main()
# Verify schedule creation was attempted
mock_create_cleanup_schedule.assert_called_once()
# Verify error was logged - check all custom_error calls
assert mock_logger.custom_error.call_count >= 1
# Find the call that contains the schedule error message
schedule_error_logged = False
for call in mock_logger.custom_error.call_args_list:
if 'Failed to configure cleanup schedule' in call[0][0]:
schedule_error_logged = True
break
assert schedule_error_logged, 'Schedule creation error should be logged'
# Verify workers were still created (startup continued)
assert mock_worker_class.call_count == 2
mock_prepare_worker.assert_not_called()
mock_start_prometheus.assert_not_called()
@patch('model_manager.worker.worker.asyncio.run')