feat: integrate PluginStore and MinIO repository into model manager activities
- Added PluginStore integration for model management. - Replaced StorageRepository with MinIORepository in Activities, Cleanup, and Training classes. - Updated training logic to handle validation files and improved data management. - Enhanced configuration for MinIO and PluginStore in connectors. - Removed deprecated model repository and storage repository files. - Updated environment variable handling for new configurations.
This commit is contained in:
@@ -20,6 +20,10 @@ def mock_env_vars():
|
||||
'PROJECT_NAME': 'test-project',
|
||||
'TRAIN_TASK_QUEUE': 'train_model-local_queue',
|
||||
'CLEANUP_TASK_QUEUE': 'cleanup-local_queue',
|
||||
'RUNTIME': 'model-manager-worker',
|
||||
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'STORE_OWNER': 'sientia',
|
||||
'STORE_REPO': 'model-library-store',
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
@@ -176,6 +180,8 @@ def test_start_prometheus_server_failure(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -191,6 +197,8 @@ async def test_main_successful_startup(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -220,6 +228,23 @@ async def test_main_successful_startup(
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
mock_activities_class.return_value = mock_activities
|
||||
|
||||
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_runtime = Mock()
|
||||
mock_runtime_class.return_value = mock_runtime
|
||||
|
||||
@@ -262,6 +287,8 @@ async def test_main_successful_startup(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -277,6 +304,8 @@ async def test_main_handles_exception(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -320,6 +349,23 @@ async def test_main_handles_exception(
|
||||
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,
|
||||
}
|
||||
|
||||
# Run main() and expect SystemExit
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
await main()
|
||||
@@ -343,6 +389,8 @@ async def test_main_handles_exception(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -358,6 +406,8 @@ async def test_main_temporal_client_configuration(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -377,6 +427,10 @@ async def test_main_temporal_client_configuration(
|
||||
'TEMPORAL_HOST': 'temporal.example.com:7233',
|
||||
'TEMPORAL_NAMESPACE': 'production',
|
||||
'TEMPORAL_USE_TLS': 'true',
|
||||
'RUNTIME': 'model-manager-worker',
|
||||
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'STORE_OWNER': 'sientia',
|
||||
'STORE_REPO': 'model-library-store',
|
||||
},
|
||||
):
|
||||
# Setup mocks
|
||||
@@ -390,6 +444,23 @@ async def test_main_temporal_client_configuration(
|
||||
mock_build_mlflow.return_value = {}
|
||||
mock_build_minio.return_value = {}
|
||||
|
||||
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_notification_handler = Mock()
|
||||
mock_notification_handler.shutdown = Mock()
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
@@ -431,6 +502,8 @@ async def test_main_temporal_client_configuration(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -446,6 +519,8 @@ async def test_main_worker_configuration(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -500,6 +575,23 @@ async def test_main_worker_configuration(
|
||||
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,
|
||||
}
|
||||
|
||||
# Run main()
|
||||
with pytest.raises(SystemExit):
|
||||
await main()
|
||||
@@ -539,6 +631,8 @@ async def test_main_worker_configuration(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -554,6 +648,8 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -603,6 +699,106 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.Worker')
|
||||
@patch('model_manager.worker.worker.client.Client')
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@patch('model_manager.worker.worker.build_minio_config')
|
||||
@patch('model_manager.worker.worker.get_logger')
|
||||
@patch('model_manager.worker.worker.start_prometheus_server')
|
||||
@patch('model_manager.worker.worker.metrics')
|
||||
async def test_main_missing_runtime_fails_fast(
|
||||
mock_metrics,
|
||||
mock_start_prometheus,
|
||||
mock_get_logger,
|
||||
mock_build_minio,
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
mock_client_class,
|
||||
mock_worker_class,
|
||||
mock_logger,
|
||||
):
|
||||
"""Test that main() fails fast when RUNTIME is missing."""
|
||||
from model_manager.worker.worker import main
|
||||
|
||||
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.shutdown = Mock()
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
|
||||
mock_activities = AsyncMock()
|
||||
mock_activities.shutdown = AsyncMock()
|
||||
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):
|
||||
await main()
|
||||
|
||||
Reference in New Issue
Block a user