chore: bump version to 2.1.0

Enhance Quality Gate Workflow for Template Management and Runtime Checks

.github/workflows/quality-gate.yml

@@ -26,4 +26,5 @@ jobs:
     with:
       project_name: 'model_manager'
       repositories: 'sientia-dataops-library,sientia-model-library'
+      runtime_check: true
     secrets: inherit

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`, 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, does **not** install the model runtime (the `sientia-module` entrypoint merges `requirements.txt` with the Plugin Store runtime and installs both before boot, SIENTIAPDE-2223; leave runtime libs unpinned in `requirements.txt` so the runtime version wins, and the quality gate's `runtime_check` fails a PR that conflicts with any warehouse runtime), sets `app_up` to 1 only after initialization, 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`.

README.md

@@ -90,7 +90,7 @@ An enterprise-grade ML model training orchestration platform built on Temporal.
 - **Scheduled Maintenance**: Automated lifecycle jobs for filesystem hygiene and stale file cleanup
 
 ### Advanced Capabilities
-- **Dynamic Runtime Provisioning**: Automated installation of required model runtimes from the Plugin Store
+- **Dynamic Runtime Provisioning**: The `sientia-module` entrypoint merges `requirements.txt` with the Plugin Store runtime and installs both before the worker starts
 - **Granular Retry Policies**: Tailored strategies for network, MLFlow, database, and filesystem operations
 - **Scalable Infrastructure**: Kubernetes-ready design with support for horizontal scaling and poller autoscaling
 - **Secure Configuration**: Environment-driven connection management with fallback to sensible defaults
@@ -148,7 +148,7 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
   - Task queue configuration and load balancing
   - Prometheus metrics server initialization
   - Notification handler setup and configuration
-  - **Runtime Installation**: Automatically installs the required model runtime from the Plugin Store
+  - **Runtime libraries**: not installed by the worker. The `sientia-module` entrypoint merges `requirements.txt` with the Plugin Store runtime (`PLUGIN_RUNTIME`), writes the audited `requirements.resolved.txt` and installs it before Python starts; on a conflict the pod does not start (SIENTIAPDE-2223). `app_up` goes to 1 only after initialization
 - **Key Features**:
   - Automatic scaling with `PollerBehaviorAutoscaling`
   - Health check endpoints for Kubernetes liveness/readiness probes
@@ -577,6 +577,18 @@ kubectl exec -n temporal <temporal-admin-tools-pod-name> -- \
     pip install -r requirements-dev.txt
     ```
 
+  4. **Install the runtime libraries**
+
+     In the cluster, the `sientia-module` entrypoint installs the runtime libraries before the
+     worker starts; the worker itself installs nothing. Locally, install them by hand from
+     `runtime/<RUNTIME>.yaml` of the Plugin Store (`Aignosi/model-plugin-warehouse`), skipping
+     `sientia_do`/`sientia_model`, which the requirements file already brings. For `basic`:
+    ```bash
+    pip install mlflow==3.8.1 pathspec==1.0.4 pandas==2.3.3 cloudpickle==3.1.2 scikit-learn==1.8.0
+    ```
+     Copy the versions from the YAML, not from this example: they change with the runtime. With
+     `RUNTIME=single` there is nothing to install.
+
 4. **Create environment configuration file**
    ```bash
    cp .env.example .env
@@ -811,6 +823,7 @@ Runs automatically on each Pull Request to `main`:
 - ✅ Tests must pass with minimum 80% coverage
 - ✅ SonarQube analysis for code quality metrics
 - ✅ Automatic version calculation based on branch pattern
+- ✅ Runtime requirements check (`runtime_check: true`): `requirements.txt` is merged with every `runtime/*.yaml` of `model-plugin-warehouse@main`, using the `sientia-module` entrypoint's own merge script; a conflict fails the gate, "runtime prevaleceu" is a warning. The warehouse changes on its own, so a push with no code change can turn red
 
 #### Deploy (`.github/workflows/deploy.yml`)
 Triggered manually via `workflow_dispatch`, with a required `chart_ref` input (no
@@ -1249,6 +1262,10 @@ sientia-dataops-model-manager/
    - Check training parameter validation errors
    - Verify input data format and required fields
 
+5. **Pod does not start, with `# CONFLICTS` in the log**
+   - The `sientia-module` entrypoint found a library whose version in `requirements.txt` conflicts with the Plugin Store runtime; the log lists each library with both specs
+   - Leave libraries the runtime already defines without a version in `requirements.txt` (the runtime's version wins), or align the version, then redeploy. The quality gate catches this on the PR against the warehouse `main`
+
 ### Debug Mode
 
 Enable debug logging by setting the log level:

docs/scenarios.md

@@ -64,3 +64,15 @@ These scenarios test the business rule validations inside `validate_train_params
 | `test_scenario_3_1_1_cleanup_with_no_temp_dirs` | Temp directory is empty. Activity completes without error. |
 | `test_scenario_3_1_2_cleanup_removes_old_temp_dirs` | Two stale directories matching `name_YYYYMMDD_HHMMSS_microseconds` are removed when older than retention. |
 | `test_scenario_3_1_3_cleanup_nonexistent_temp_path` | Target path does not exist. Handled gracefully without error. |
+
+## 4. PluginStore Wiring (`test_plugin_store_wiring.py`)
+
+`Activities` receives `plugin_store_config` and builds the `PluginStore` itself. These scenarios run `01-linear-regression-basic.json` on a worker built from a per-test config (`activities_factory`) against the Gitea `model-store` repo.
+
+| Test Function | Config Change | Expected Status | Description |
+|---|---|---|---|
+| `test_scenario_4_1_1_plugin_store_from_env_config` | `STORE_*` env vars → `build_plugin_store_config()` | `TRAINING_SUCCESS` | Same path as the worker. With `STORE_BRANCH` unset the store resolves the branch to `main`. |
+| `test_scenario_4_1_2_explicit_branch` | `branch = 'main'` | `TRAINING_SUCCESS` | Explicit branch skips resolution and is used for every read. |
+| `test_scenario_4_1_3_explicit_branch_without_model` | `branch = 'e2e-no-linear-regression'` | `ORCHESTRATOR_VALIDATION_ERROR` | Branch forked from `main` without `models/linear_regression/index.yaml`; error mentions `linear_regression/index.yaml`. |
+| `test_scenario_4_1_4_nonexistent_repository` | `repo = 'no-such-store'` | `ORCHESTRATOR_VALIDATION_ERROR` | `Activities` still builds (branch falls back to `main`); fails when loading model metadata. |
+| `test_scenario_4_1_5_wrong_credentials` | `password = 'wrong-password'` | `ORCHESTRATOR_VALIDATION_ERROR` | Gitea rejects the read; error mentions `401`. |

model_manager/utils/repository/data_manager_repository.py

@@ -357,7 +357,8 @@ class DataManagerRepository(SientiaMonitoring):
         This helper mirrors the previous TrainingRepository.after_train_calculation
         behavior, assuming that predictions (y_pred/y_train_pred) are already on the
         correct scale for metric calculation (any scaling is handled inside the
-        model wrapper).
+        model wrapper). Pairs with a NaN target or a NaN prediction are left out of
+        mse/mae/r2. With a single valid pair left, r2_val is NaN (R² needs 2 samples).
 
         Args:
             tmr: Training result containing:
@@ -368,6 +369,10 @@ class DataManagerRepository(SientiaMonitoring):
 
         Return:
             TrainModelResult: Same object with mse_val, mae_val and r2_val set.
+
+        Raises:
+            ValueError: If y_pred is not set, val_data and y_pred share no indices,
+                or no pair has both a valid target and a valid prediction.
         """
         if tmr.y_pred is None:
             raise ValueError('y_pred must be set before computing regression metrics')
@@ -405,6 +410,24 @@ class DataManagerRepository(SientiaMonitoring):
         y_true_val = y_true_val.loc[common_index]
         y_pred_val = y_pred_val.loc[common_index]
 
+        # NaN target, or NaN prediction for rows with NaN features, carries no error signal
+        # (SIENTIAPDE-2215).
+        valid = y_true_val.notna() & y_pred_val.notna()
+        if not valid.any():
+            raise ValueError(
+                'No validation rows with both target and prediction: all '
+                f'{len(valid)} overlapping rows have NaN'
+            )
+        if not valid.all():
+            n_valid = int(valid.sum())
+            message = (
+                f'Ignoring {len(valid) - n_valid} of {len(valid)} validation row(s) with NaN '
+                f'target or prediction; regression metrics use the remaining {n_valid}'
+            )
+            self.warning(message, metadata)
+        y_true_val = y_true_val[valid]
+        y_pred_val = y_pred_val[valid]
+
         # Metrics helpers already round to 2 decimals.
         tmr.mse_val = mse(y_true_val, y_pred_val)
         tmr.mae_val = mae(y_true_val, y_pred_val)

model_manager/worker/worker.py

@@ -83,7 +83,9 @@ 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
+    3. Initializes notification handler and Plugin Store client, then sets ``app_up``
+       (what the probes check) to 1. The runtime libraries are already installed: the
+       sientia-module entrypoint merges them into the resolved requirements before boot
     4. Creates and configures activities
     5. Starts Temporal client and workers
     6. Manages worker lifecycle and graceful shutdown
@@ -125,8 +127,6 @@ 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'],
@@ -144,7 +144,8 @@ async def main():
         metrics_controller=metrics_controller,
     )
 
-    await plugin_store.install_runtime(runtime_name=runtime)
+    # The probes only check app_up, so it goes to 1 only after initialization.
+    metrics.APP_UP.labels(pod_id=POD_ID).set(1)
 
     activities = Activities(
         postgres_config=build_postgres_config(),
@@ -232,8 +233,8 @@ def start_prometheus_server(logger: SientiaLogger, metadata: dict[str, str | Non
     """
     Starts the Prometheus metrics server for monitoring and observability.
 
-    This function initializes the Prometheus HTTP server on the configured port
-    and sets the application health metric to indicate the service is running.
+    This function initializes the Prometheus HTTP server on the configured port.
+    ``main`` sets ``app_up`` to 1 later, once initialization is done.
 
     The server exposes metrics that can be scraped by Prometheus for monitoring
     the health and performance of the Model Manager worker.
@@ -249,7 +250,6 @@ def start_prometheus_server(logger: SientiaLogger, metadata: dict[str, str | Non
         port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
         logger.custom_info(f'Prometheus server initialized on port {port}.', metadata)
-        metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:  # noqa: BLE001
         logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
         os._exit(1)

model_manager/workflows/train_model.py

@@ -131,7 +131,9 @@ class TrainModel:
             training_succeeded = True
         finally:
             try:
-                if train_result is not None:
+                # train_result is only None when training raised, so the false branch
+                # always leaves through the re-raise and never reaches the return below.
+                if train_result is not None:  # pragma: no branch
                     await self._cleanup_resources(
                         run_dir=train_result.get('run_dir'),
                         metadata=metadata,

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "model-manager"
-version = "2.0.3"
+version = "2.1.0"
 description = "Sientia DataOps Model Manager - ML Model Orchestration System"
 readme = "README.md"
 requires-python = ">=3.11"

requirements-local.txt

@@ -1,11 +1,19 @@
-temporalio==1.23.0
-psycopg2-binary==2.9.11
-sqlalchemy==2.0.49
-boto3==1.42.70
-botocore==1.42.70
+temporalio
+psycopg2-binary
+# SQLAlchemy 2.1 makes psycopg (v3) the default driver for postgresql:// URLs; sientia_do 1.12.2
+# builds that URL and ships only psycopg2-binary, so 2.1 fails with "No module named psycopg".
+sqlalchemy<2.1
+boto3
+botocore
 git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.12.2
 git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.14.0
-prometheus-client==0.23.1
-beautifulsoup4==4.12.3
-evidently==0.6.7
-jsonschema==4.26.0
+prometheus-client
+beautifulsoup4
+# reports.py uses the Report/metric_preset API that 0.7 moved to evidently.legacy.
+evidently<0.7
+jsonschema
+# Imported directly by the worker, without a version: the sientia-module entrypoint merges this
+# file with the Plugin Store runtime, and the runtime's version wins (SIENTIAPDE-2223).
+mlflow
+pandas
+scikit-learn

requirements.txt

@@ -1,11 +1,19 @@
-temporalio==1.23.0
-psycopg2-binary==2.9.11
-sqlalchemy==2.0.50
-boto3==1.42.70
-botocore==1.42.70
+temporalio
+psycopg2-binary
+# SQLAlchemy 2.1 makes psycopg (v3) the default driver for postgresql:// URLs; sientia_do 1.12.2
+# builds that URL and ships only psycopg2-binary, so 2.1 fails with "No module named psycopg".
+sqlalchemy<2.1
+boto3
+botocore
 sientia_do>=1.12.2
-sientia_model>=0.8.1
-prometheus-client==0.23.1
-beautifulsoup4==4.12.3
-evidently==0.6.7
-jsonschema==4.26.0
+sientia_model>=0.14.0
+prometheus-client
+beautifulsoup4
+# reports.py uses the Report/metric_preset API that 0.7 moved to evidently.legacy.
+evidently<0.7
+jsonschema
+# Imported directly by the worker, without a version: the sientia-module entrypoint merges this
+# file with the Plugin Store runtime, and the runtime's version wins (SIENTIAPDE-2223).
+mlflow
+pandas
+scikit-learn

tests/utils/repository/test_data_manager_repository.py

@@ -92,6 +92,53 @@ def test_ensure_date_column_parsed_invalid_raises():
         dmr._ensure_date_column_parsed(df, p)
 
 
+def test_ensure_date_column_parsed_unmapped_format_raises():
+    df = pd.DataFrame({'ts': ['2024-01-01 10:00:00']})
+    p = _params(date_column='ts')
+    # Bypasses validate_business_rules, which would reject it earlier
+    p.date_format = 'dd/MM/yyyy HH:mm'
+    with pytest.raises(ValueError, match='is not mapped to a strftime pattern'):
+        dmr._ensure_date_column_parsed(df, p)
+
+
+def test_drop_rows_with_missing_timestamp_without_date_column_returns_input():
+    repo = dmr.DataManagerRepository(MagicMock())
+    df = pd.DataFrame({'v1': [1.0, 2.0]})
+    assert repo._drop_rows_with_missing_timestamp(df, _params(), {}) is df
+
+
+def test_drop_rows_with_missing_timestamp_non_object_dtype_drops_nat():
+    repo = dmr.DataManagerRepository(MagicMock())
+    df = pd.DataFrame(
+        {'timestamp': pd.to_datetime(['2024-01-01', None, '2024-01-03']), 'v1': [1, 2, 3]}
+    )
+    out = repo._drop_rows_with_missing_timestamp(df, _params(), {})
+    assert list(out['v1']) == [1, 3]
+
+
+def test_coerce_non_timestamp_columns_to_numeric_skips_date_column():
+    repo = dmr.DataManagerRepository(MagicMock())
+    df = pd.DataFrame({'timestamp': ['2024-01-01 10:00:00'], 'v1': ['1.5']})
+    out = repo._coerce_non_timestamp_columns_to_numeric(df, _params(), {})
+    assert out['timestamp'].iloc[0] == '2024-01-01 10:00:00'
+    assert float(out['v1'].iloc[0]) == 1.5
+
+
+def test_set_timezone_on_index_converts_aware_index_to_utc():
+    repo = dmr.DataManagerRepository(MagicMock())
+    idx = pd.DatetimeIndex(['2024-01-01 00:00'], tz='America/Sao_Paulo')
+    out = repo._set_timezone_on_index(pd.DataFrame({'v1': [1]}, index=idx))
+    assert str(out.index.tz) == 'UTC'
+    assert out.index[0] == pd.Timestamp('2024-01-01 03:00', tz='UTC')
+
+
+def test_set_timezone_on_index_non_datetime_index_raises():
+    repo = dmr.DataManagerRepository(MagicMock())
+    df = pd.DataFrame({'v1': [1]})
+    with pytest.raises(ValueError, match='Index is not a DatetimeIndex'):
+        repo._set_timezone_on_index(df)
+
+
 def test_prepare_training_data_csv_load_failure():
     repo = dmr.DataManagerRepository(MagicMock())
     p = _params()
@@ -309,6 +356,46 @@ def test_compute_regression_metrics_no_overlap():
         repo.compute_regression_metrics(tmr, MagicMock())
 
 
+def test_compute_regression_metrics_ignores_nan_pairs():
+    """NaN target or NaN prediction rows are left out of mse/mae/r2 (SIENTIAPDE-2215)."""
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    idx = pd.Index([0, 1, 2, 3])
+    tmr = TrainModelResult(
+        params=p,
+        train_data=pd.DataFrame({'t': [1.0, 2.0, 3.0, 4.0]}, index=idx),
+        val_data=pd.DataFrame({'t': [1.0, np.nan, 3.0, 4.0]}, index=idx),
+        y_pred=pd.DataFrame({'p': [1.0, 2.0, np.nan, 4.0]}, index=idx),
+    )
+
+    wrapper = MagicMock()
+    wrapper.model = object()  # no sklearn .regr -> equation extraction is skipped
+    with patch.object(repo, 'warning') as warn:
+        out = repo.compute_regression_metrics(tmr, wrapper)
+
+    assert out.mse_val == pytest.approx(0.0)
+    assert out.mae_val == pytest.approx(0.0)
+    assert out.r2_val == pytest.approx(1.0)
+    warn.assert_called_once()
+    assert 'Ignoring 2 of 4' in warn.call_args.args[0]
+
+
+def test_compute_regression_metrics_all_nan_pairs_raises():
+    repo = dmr.DataManagerRepository(MagicMock())
+    p = TrainModelParams.from_dict(_minimal_dict_for_prepare())
+    idx = pd.Index([0, 1])
+    tmr = TrainModelResult(
+        params=p,
+        train_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
+        val_data=pd.DataFrame({'t': [1.0, 2.0]}, index=idx),
+        y_pred=pd.DataFrame({'p': [np.nan, np.nan]}, index=idx),
+    )
+    wrapper = MagicMock()
+
+    with pytest.raises(ValueError, match='No validation rows with both target and prediction'):
+        repo.compute_regression_metrics(tmr, wrapper)
+
+
 def test_compute_regression_metrics_linear_equation():
     repo = dmr.DataManagerRepository(MagicMock())
     p = TrainModelParams.from_dict(_minimal_dict_for_prepare())

tests/worker/test_worker.py

@@ -132,9 +132,8 @@ def test_start_prometheus_server_success(
     # Verify HTTP server started
     mock_start_http_server.assert_called_once_with(9090)
 
-    # Verify APP_UP metric was set to 1
-    mock_metrics.APP_UP.labels.assert_called_once_with(pod_id='test-pod-123')
-    mock_app_up.set.assert_called_once_with(1)
+    # app_up goes to 1 in main(), after initialization.
+    mock_app_up.set.assert_not_called()
     mock_logger.custom_info.assert_called_once()
 
 
@@ -233,9 +232,6 @@ async def test_main_successful_startup(
     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',
@@ -284,7 +280,10 @@ async def test_main_successful_startup(
     # Verify cleanup was performed
     mock_notification_handler.shutdown.assert_called_once()
     mock_activities.shutdown.assert_called_once()
-    mock_app_up.set.assert_called_with(0)
+    # Ready after initialization, down again on shutdown.
+    assert [c.args[0] for c in mock_app_up.set.call_args_list] == [1, 0]
+    # The runtime libraries come from the sientia-module entrypoint, not from the worker.
+    mock_plugin_store_instance.install_runtime.assert_not_called()
 
 
 @pytest.mark.asyncio
@@ -360,9 +359,6 @@ async def test_main_handles_exception(
     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',
@@ -459,9 +455,6 @@ async def test_main_temporal_client_configuration(
         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',
@@ -594,9 +587,6 @@ async def test_main_worker_configuration(
     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',
@@ -716,9 +706,6 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
     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',
@@ -816,9 +803,6 @@ async def test_main_missing_runtime_uses_single_fallback(
     mock_prepare_worker.return_value = mock_worker_instance
 
     mock_plugin_store_instance = AsyncMock()
-    mock_plugin_store_instance.install_runtime = AsyncMock(
-        return_value={'runtime': 'single', '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',
@@ -848,12 +832,10 @@ async def test_main_missing_runtime_uses_single_fallback(
 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)