From 0f5e4a5ceffc22f81530cfe7b00a32b9326a9054 Mon Sep 17 00:00:00 2001 From: PedroHMCosme Date: Thu, 22 Oct 2026 11:34:04 -0300 Subject: [PATCH] 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. --- CLAUDE.md | 2 +- README.md | 1 - model_manager/worker/worker.py | 66 ++++++--- requirements.txt | 11 ++ tests/worker/test_worker.py | 240 ++++++++++----------------------- 5 files changed, 127 insertions(+), 193 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 6935bd8..3ced319 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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`. diff --git a/README.md b/README.md index 5234dca..4483917 100644 --- a/README.md +++ b/README.md @@ -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--queue`, `cleanup_files--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 | diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index c2bad6b..70bdb32 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -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) - 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', + loaded_before_install = _loaded_distribution_versions() + await plugin_store.install_runtime(runtime_name=runtime) + + 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(), diff --git a/requirements.txt b/requirements.txt index 46894ca..9ae8fa6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 82df3f6..9fb43ff 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -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