fix(worker): pin the runtime versions and refuse to boot when install_runtime swaps an imported package (SIENTIAPDE-2214)
Replaces the re-exec. requirements.txt now pins the union of the basic and xgboost runtime libraries, so install_runtime has nothing to change in the packages the worker already imported. If a runtime pin changes and requirements.txt is not updated, the worker logs which packages changed and exits before starting the workers. test_main_entrypoint assigned worker.main = AsyncMock() without restoring it, so every later test of main() ran the mock; it now uses patch.object.
This commit is contained in:
PedroHMCosme
@@ -3,7 +3,7 @@
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -30,16 +30,6 @@ def mock_env_vars():
|
||||
yield env_vars
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def runtime_ready(monkeypatch):
|
||||
"""Boot as the re-executed process by default; the real os.execv must never run in tests."""
|
||||
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '1')
|
||||
|
||||
|
||||
class _StopBootError(Exception):
|
||||
"""Stops main() at a chosen point, e.g. where os.execv would replace the process."""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger."""
|
||||
@@ -854,172 +844,14 @@ async def test_main_missing_runtime_uses_single_fallback(
|
||||
assert mock_prepare_worker.call_args_list[1].kwargs['runtime'] == 'single'
|
||||
|
||||
|
||||
_PLUGIN_STORE_CONFIG = {
|
||||
'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,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
||||
# Stops boot at Activities, so the red phase does not try to reach a real Temporal server.
|
||||
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
||||
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
||||
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
||||
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
||||
@patch(
|
||||
'model_manager.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
||||
)
|
||||
@patch('model_manager.worker.worker.get_logger')
|
||||
@patch('model_manager.worker.worker.start_prometheus_server')
|
||||
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
||||
async def test_main_reexecs_after_installing_runtime(
|
||||
_mock_dirs,
|
||||
_mock_prometheus,
|
||||
_mock_get_logger,
|
||||
_mock_mongo,
|
||||
_mock_minio,
|
||||
_mock_mlflow,
|
||||
_mock_postgres,
|
||||
_mock_ps_config,
|
||||
mock_plugin_store_class,
|
||||
_mock_notification,
|
||||
mock_activities_class,
|
||||
monkeypatch,
|
||||
):
|
||||
"""First boot installs the runtime, then replaces the process so imports see it on disk."""
|
||||
from model_manager.worker import worker as worker_module
|
||||
|
||||
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '0')
|
||||
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
||||
|
||||
with (
|
||||
patch('model_manager.worker.worker.os.execv', side_effect=_StopBootError) as m_execv,
|
||||
patch.object(sys, 'orig_argv', ['python', '-m', 'model_manager.worker.worker']),
|
||||
pytest.raises(_StopBootError),
|
||||
):
|
||||
await worker_module.main()
|
||||
|
||||
mock_plugin_store_class.return_value.install_runtime.assert_awaited_once_with(
|
||||
runtime_name='basic'
|
||||
)
|
||||
assert os.environ['SIENTIA_RUNTIME_READY'] == '1'
|
||||
m_execv.assert_called_once_with(
|
||||
sys.executable, [sys.executable, '-m', 'model_manager.worker.worker']
|
||||
)
|
||||
mock_activities_class.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
||||
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
||||
@patch('model_manager.worker.worker.build_postgres_config', return_value={})
|
||||
@patch('model_manager.worker.worker.build_mlflow_config', return_value={})
|
||||
@patch('model_manager.worker.worker.build_minio_config', return_value={})
|
||||
@patch(
|
||||
'model_manager.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
||||
)
|
||||
@patch('model_manager.worker.worker.get_logger')
|
||||
@patch('model_manager.worker.worker.start_prometheus_server')
|
||||
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
||||
async def test_main_skips_install_when_runtime_ready(
|
||||
_mock_dirs,
|
||||
_mock_prometheus,
|
||||
_mock_get_logger,
|
||||
_mock_mongo,
|
||||
_mock_minio,
|
||||
_mock_mlflow,
|
||||
_mock_postgres,
|
||||
_mock_ps_config,
|
||||
mock_plugin_store_class,
|
||||
_mock_notification,
|
||||
mock_activities_class,
|
||||
):
|
||||
"""The re-executed process skips the install and the exec and goes on booting."""
|
||||
from model_manager.worker import worker as worker_module
|
||||
|
||||
mock_plugin_store_class.return_value.install_runtime = AsyncMock(return_value=None)
|
||||
|
||||
with (
|
||||
patch('model_manager.worker.worker.os.execv') as m_execv,
|
||||
pytest.raises(_StopBootError),
|
||||
):
|
||||
await worker_module.main()
|
||||
|
||||
mock_plugin_store_class.return_value.install_runtime.assert_not_awaited()
|
||||
m_execv.assert_not_called()
|
||||
mock_activities_class.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.RUNTIME', 'basic')
|
||||
@patch('model_manager.worker.worker.Activities', side_effect=_StopBootError)
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config', return_value=_PLUGIN_STORE_CONFIG)
|
||||
@patch(
|
||||
'model_manager.worker.worker.build_mongodb_config',
|
||||
return_value={'connection_string': 'mongodb://test', 'database_name': 'db', 'uri': 'u'},
|
||||
)
|
||||
@patch('model_manager.worker.worker.get_logger')
|
||||
@patch('model_manager.worker.worker.start_prometheus_server')
|
||||
@patch('model_manager.worker.worker.ensure_runtime_directories')
|
||||
async def test_main_install_failure_does_not_reexec(
|
||||
_mock_dirs,
|
||||
_mock_prometheus,
|
||||
_mock_get_logger,
|
||||
_mock_mongo,
|
||||
_mock_ps_config,
|
||||
mock_plugin_store_class,
|
||||
_mock_notification,
|
||||
mock_activities_class,
|
||||
monkeypatch,
|
||||
):
|
||||
"""A failed install propagates without marking the runtime ready or replacing the process."""
|
||||
from model_manager.worker import worker as worker_module
|
||||
|
||||
monkeypatch.setenv('SIENTIA_RUNTIME_READY', '0')
|
||||
mock_plugin_store_class.return_value.install_runtime = AsyncMock(
|
||||
side_effect=RuntimeError('pip failed')
|
||||
)
|
||||
|
||||
with (
|
||||
patch('model_manager.worker.worker.os.execv') as m_execv,
|
||||
pytest.raises(RuntimeError, match='pip failed'),
|
||||
):
|
||||
await worker_module.main()
|
||||
|
||||
assert os.environ['SIENTIA_RUNTIME_READY'] == '0'
|
||||
m_execv.assert_not_called()
|
||||
mock_activities_class.assert_not_called()
|
||||
|
||||
|
||||
@patch('model_manager.worker.worker.asyncio.run')
|
||||
def test_main_entrypoint(mock_asyncio_run):
|
||||
"""Test the __main__ entrypoint."""
|
||||
# Import and execute the main block
|
||||
with patch.object(sys, 'argv', ['worker.py']):
|
||||
import model_manager.worker.worker as worker_module
|
||||
|
||||
# Simulate running the module
|
||||
worker_module.main = AsyncMock()
|
||||
import model_manager.worker.worker as worker_module
|
||||
|
||||
# patch.object restores main afterwards; assigning it leaked the mock into later tests.
|
||||
with patch.object(sys, 'argv', ['worker.py']), patch.object(worker_module, 'main', AsyncMock()):
|
||||
# This would normally be called by asyncio.run(main())
|
||||
# We just verify the pattern is correct
|
||||
assert callable(worker_module.main)
|
||||
@@ -1071,3 +903,67 @@ def test_start_prometheus_server_prints_failure(
|
||||
# Agora o erro é logado via logger crítico
|
||||
mock_logger.custom_critical.assert_called_once()
|
||||
mock_exit.assert_called_once_with(1)
|
||||
|
||||
|
||||
def test_loaded_distribution_versions_reads_imported_packages():
|
||||
"""An imported module maps to its distribution's installed version."""
|
||||
import importlib.metadata
|
||||
|
||||
from model_manager.worker.worker import _loaded_distribution_versions
|
||||
|
||||
assert _loaded_distribution_versions()['pytest'] == importlib.metadata.version('pytest')
|
||||
|
||||
|
||||
def test_runtime_changed_loaded_packages_reports_only_changes():
|
||||
"""Only distributions whose version on disk differs after the install are reported."""
|
||||
import importlib.metadata
|
||||
|
||||
from model_manager.worker.worker import _runtime_changed_loaded_packages
|
||||
|
||||
on_disk = {'scikit-learn': '1.8.0', 'numpy': '1.26.4'}
|
||||
|
||||
def fake_version(dist):
|
||||
if dist not in on_disk:
|
||||
raise importlib.metadata.PackageNotFoundError(dist)
|
||||
return on_disk[dist]
|
||||
|
||||
before = {'scikit-learn': '1.9.1', 'numpy': '1.26.4', 'gone': '1.0'}
|
||||
with patch('model_manager.worker.worker.importlib.metadata.version', side_effect=fake_version):
|
||||
changed = _runtime_changed_loaded_packages(before)
|
||||
|
||||
assert changed == {'scikit-learn': '1.9.1 -> 1.8.0', 'gone': '1.0 -> removed'}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_main_runtime_changed_loaded_package_exits(mock_env_vars, mock_logger):
|
||||
"""Boot stops before Activities when install_runtime changed an imported package."""
|
||||
import model_manager.worker.worker as worker_module
|
||||
|
||||
plugin_store = Mock()
|
||||
plugin_store.install_runtime = AsyncMock(return_value=None)
|
||||
mongo = {'connection_string': 'mongodb://test', 'database_name': 'test_db', 'uri': 'x'}
|
||||
with (
|
||||
patch.object(worker_module, 'RUNTIME', 'basic'),
|
||||
patch.object(worker_module, 'ensure_runtime_directories'),
|
||||
patch.object(worker_module, 'metrics'),
|
||||
patch.object(worker_module, 'start_prometheus_server'),
|
||||
patch.object(worker_module, 'get_logger', return_value=mock_logger),
|
||||
patch.object(worker_module, 'build_mongodb_config', return_value=mongo),
|
||||
patch.object(worker_module, 'NotificationHandler'),
|
||||
patch.object(worker_module, 'build_plugin_store_config', return_value=MagicMock()),
|
||||
patch.object(worker_module, 'PluginStore', return_value=plugin_store),
|
||||
patch.object(
|
||||
worker_module,
|
||||
'_runtime_changed_loaded_packages',
|
||||
return_value={'scikit-learn': '1.9.1 -> 1.8.0'},
|
||||
),
|
||||
patch.object(worker_module, 'Activities') as mock_activities_class,
|
||||
patch.object(worker_module.sys, 'exit', side_effect=SystemExit(1)) as mock_exit,
|
||||
):
|
||||
with pytest.raises(SystemExit):
|
||||
await worker_module.main()
|
||||
|
||||
mock_exit.assert_called_once_with(1)
|
||||
mock_activities_class.assert_not_called()
|
||||
message = mock_logger.custom_critical.call_args.args[0]
|
||||
assert "'scikit-learn': '1.9.1 -> 1.8.0'" in message
|
||||
|
||||
Reference in New Issue
Block a user