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
@@ -45,7 +45,7 @@ Coverage gate in CI is 80% minimum; this repo targets 100% in practice — don't
|
||||
## Architecture
|
||||
|
||||
### Temporal layering
|
||||
- **Worker** (`model_manager/worker/worker.py`): entrypoint. Builds config from env, connects Postgres/MLFlow/MinIO/Plugin Store, installs the `RUNTIME` model runtime via `PluginStore.install_runtime` and then re-executes itself once (`os.execv`, flagged by the internal `SIENTIA_RUNTIME_READY=1`) so every import loads the runtime's package versions (SIENTIAPDE-2214), starts two Temporal workers (one per task queue) plus Prometheus metrics server and the cleanup schedule. All heavy imports are wrapped in `workflow.unsafe.imports_passed_through()` — required because Temporal workflow sandboxing forbids non-deterministic imports at module scope outside that block.
|
||||
- **Worker** (`model_manager/worker/worker.py`): entrypoint. Builds config from env, connects Postgres/MLFlow/MinIO/Plugin Store, installs the `RUNTIME` model runtime via `PluginStore.install_runtime`, starts two Temporal workers (one per task queue) plus Prometheus metrics server and the cleanup schedule. All heavy imports are wrapped in `workflow.unsafe.imports_passed_through()` — required because Temporal workflow sandboxing forbids non-deterministic imports at module scope outside that block.
|
||||
- **Workflows** (`model_manager/workflows/`): pure orchestration, no I/O beyond `execute_activity_method` calls.
|
||||
- `train_model.py` (`TrainModel`, workflow name `train_model`): validate `experiment_run_id` → `load_model_metadata` → `validate_train_params` → update DB status `ORCHESTRATOR_WAITING_PROC` → `train_model` activity → update DB status `TRAINING_SUCCESS`/`TRAINING_ERROR` → `cleanup_resources` (always runs in `finally` if a `run_dir` exists, even on failure). Uses distinct `RetryPolicy`s per operation class: `network_retry_policy` (transient I/O), `no_retry_policy` (training/validation — data errors are permanent, and retrying would duplicate MLflow writes), `database_retry_policy` (Postgres). Timeouts are env-driven (`TIMEOUT_*`), extend `TIMEOUT_TRAIN_MODEL` instead of adding retries for slow training.
|
||||
- `cleanup_files.py` (`CleanupFiles`): scheduled (`cleanup_schedule.py`, default daily at midnight UTC via `CLEANUP_CRON`) deletion of stale local temp dirs matching `name_YYYYMMDD_HHMMSS_microseconds`, older than `CLEANUP_RETENTION_HOURS`.
|
||||
|
||||
@@ -990,7 +990,6 @@ Values are read in `model_manager/utils/connectors_config.py` and `model_manager
|
||||
| `TEMPORAL_NAMESPACE` | Temporal namespace | `model-manager` | No |
|
||||
| `TEMPORAL_USE_TLS` | Use TLS for Temporal gRPC (`true`/`false`). Set `true` when the endpoint serves TLS or you get HTTP redirects (for example 308) to HTTPS | `false` | No |
|
||||
| `RUNTIME` | Suffix for worker task queues (`train_model-<runtime>-queue`, `cleanup_files-<runtime>-queue`) | `single` (via `_get_runtime`) | No |
|
||||
| `SIENTIA_RUNTIME_READY` | Internal: set by the worker itself before re-executing after `install_runtime` (SIENTIAPDE-2214). **Never set it in `.env` or the chart**: it skips the runtime install | unset | No |
|
||||
| `TRAIN_TASK_QUEUE` | Used by **clients** (for example `scripts/run_training_test.py`), not by the worker process | unset | No |
|
||||
| `CLEANUP_TASK_QUEUE` | Used by **clients** (for example `scripts/run_cleanup_test.py`), not by the worker | unset | No |
|
||||
| `POSTGRES_HOST` | PostgreSQL hostname | `localhost` | Yes |
|
||||
|
||||
@@ -21,8 +21,6 @@ Environment Variables:
|
||||
- TEMPORAL_NAMESPACE: Temporal namespace (default: model-manager)
|
||||
- TEMPORAL_USE_TLS: Enable TLS for Temporal connection (default: false)
|
||||
- RUNTIME: Runtime identifier used in queue naming (default: single)
|
||||
- SIENTIA_RUNTIME_READY: Internal, set by the worker itself before re-executing after
|
||||
install_runtime. Never set it in the deployment: it skips the runtime install.
|
||||
- POD_ID: Kubernetes pod identifier for metrics
|
||||
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
|
||||
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
|
||||
@@ -34,6 +32,7 @@ from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import asyncio
|
||||
import importlib.metadata
|
||||
import os
|
||||
import sys
|
||||
|
||||
@@ -78,6 +77,35 @@ def _get_runtime(runtime: str | None) -> str:
|
||||
return normalized_runtime or 'single'
|
||||
|
||||
|
||||
def _loaded_distribution_versions() -> dict[str, str]:
|
||||
"""Version on disk of every installed distribution whose top-level module is imported."""
|
||||
loaded = {name.partition('.')[0] for name in sys.modules}
|
||||
return {
|
||||
dist: importlib.metadata.version(dist)
|
||||
for module, dists in importlib.metadata.packages_distributions().items()
|
||||
if module in loaded
|
||||
for dist in dists
|
||||
}
|
||||
|
||||
|
||||
def _runtime_changed_loaded_packages(before: dict[str, str]) -> dict[str, str]:
|
||||
"""
|
||||
Packages that were imported before ``install_runtime`` and whose version changed on disk.
|
||||
|
||||
The process keeps the old code in memory, so a later lazy import mixes two versions
|
||||
(SIENTIAPDE-2214: sklearn 1.9.1 in memory, 1.8.0 on disk, ``_cyutility`` import error).
|
||||
"""
|
||||
changed = {}
|
||||
for dist, old in before.items():
|
||||
try:
|
||||
new = importlib.metadata.version(dist)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
new = 'removed'
|
||||
if new != old:
|
||||
changed[dist] = f'{old} -> {new}'
|
||||
return changed
|
||||
|
||||
|
||||
async def main():
|
||||
"""
|
||||
Main entry point for the Model Manager worker application.
|
||||
@@ -85,12 +113,11 @@ async def main():
|
||||
This function initializes and starts all components of the worker:
|
||||
1. Sets up logging and metadata
|
||||
2. Starts Prometheus metrics server
|
||||
3. Initializes notification handler
|
||||
4. Installs the PluginStore runtime and re-executes itself once (`SIENTIA_RUNTIME_READY=1`)
|
||||
so every import loads the runtime's package versions
|
||||
5. Creates and configures activities
|
||||
6. Starts Temporal client and workers
|
||||
7. Manages worker lifecycle and graceful shutdown
|
||||
3. Initializes notification handler, then installs the runtime; exits if the install
|
||||
changed the version of a package the worker already imported
|
||||
4. Creates and configures activities
|
||||
5. Starts Temporal client and workers
|
||||
6. Manages worker lifecycle and graceful shutdown
|
||||
|
||||
The function runs indefinitely until interrupted or an error occurs.
|
||||
On error, it performs cleanup and exits with a non-zero status code.
|
||||
@@ -129,6 +156,8 @@ async def main():
|
||||
|
||||
metrics_controller = MetricsController(logger=logger)
|
||||
|
||||
logger.custom_info(f'Installing runtime {runtime}', metadata)
|
||||
|
||||
plugin_store_parameters = build_plugin_store_config()
|
||||
plugin_store = PluginStore(
|
||||
base_url=plugin_store_parameters['base_url'],
|
||||
@@ -146,19 +175,18 @@ async def main():
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
if os.environ.get('SIENTIA_RUNTIME_READY') != '1':
|
||||
logger.custom_info(f'Installing runtime {runtime}', metadata)
|
||||
loaded_before_install = _loaded_distribution_versions()
|
||||
await plugin_store.install_runtime(runtime_name=runtime)
|
||||
os.environ['SIENTIA_RUNTIME_READY'] = '1'
|
||||
logger.custom_info(f'Runtime {runtime} installed, re-executing worker', metadata)
|
||||
# ponytail: fresh interpreter so loaded modules match the runtime on disk (SIENTIAPDE-2214).
|
||||
# The install_runtime success metrics (count/lag) are lost to the exec; no dashboard uses them.
|
||||
os.execv(sys.executable, [sys.executable, *sys.orig_argv[1:]]) # noqa: S606
|
||||
else:
|
||||
logger.custom_info(
|
||||
f'SIENTIA_RUNTIME_READY=1, runtime {runtime} already installed, skipping install',
|
||||
|
||||
changed = _runtime_changed_loaded_packages(loaded_before_install)
|
||||
if changed:
|
||||
logger.custom_critical(
|
||||
f'Runtime {runtime} changed packages already imported by the worker: {changed}. '
|
||||
'Pin them in requirements.txt to the runtime versions.',
|
||||
metadata,
|
||||
)
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
||||
sys.exit(1)
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=build_postgres_config(),
|
||||
|
||||
@@ -9,3 +9,14 @@ prometheus-client==0.23.1
|
||||
beautifulsoup4==4.12.3
|
||||
evidently==0.6.7
|
||||
jsonschema==4.26.0
|
||||
# Same versions as the plugin runtimes (model-plugin-factory runtime/basic.yaml and xgboost.yaml).
|
||||
# The worker imports these before install_runtime; if the runtime installed a different version,
|
||||
# the old code would stay in memory, so the worker refuses to boot (SIENTIAPDE-2214).
|
||||
# Changing a runtime pin means changing it here too.
|
||||
mlflow==3.8.1
|
||||
pathspec==1.0.4
|
||||
pandas==2.3.3
|
||||
cloudpickle==3.1.2
|
||||
scikit-learn==1.8.0
|
||||
joblib==1.5.3
|
||||
rich==14.3.2
|
||||
|
||||
@@ -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()
|
||||
|
||||
# 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