sientia-dataops-model-manager/model-manager-plugin-store-migration-plan.md

---
tags:
  - engineering
  - sientia
  - runtime-system
  - model-manager
  - plugin-store
  - migration-plan
created: 2026-03-02
modified: 2026-03-02
created_by: Vitor Pimentel
modified_by: Vitor Pimentel
status: draft
---

# Sientia Model Manager — PluginStore Migration Plan

> Implementation plan for migrating `sientia-dataops-model-manager` to use the Sientia PluginStore for runtime installation and model retrieval, aligned with the runtime architecture described in `analytics.md`.

## Summary

1. [[#Objectives and Scope|Objectives and Scope]] — What this migration must achieve
2. [[#Existing State Overview (model-manager)|Existing State Overview]] — Current responsibilities and coupling points
3. [[#Requirements Mapping|Requirements Mapping]] — Functional and non-functional requirements
4. [[#Target Architecture|Target Architecture]] — Desired runtime and model-loading architecture
5. [[#Implementation Plan|Implementation Plan]] — Phased, detailed changes to apply
6. [[#Testing Strategy|Testing Strategy]] — How to validate the new behavior
7. [[#Rollout and Migration Strategy|Rollout and Migration Strategy]] — How to safely roll out and deprecate old paths
8. [[#Potential Model Library Changes|Potential Model Library Changes]] — Expected impact on `sientia-model-library`
9. [[#Related Documents|Related Documents]] — Cross-links to supporting documents

---

## Objectives and Scope

The goal of this work is to evolve `sientia-dataops-model-manager` so that:

- It **detects the runtime** before the worker starts, using the `RUNTIME` environment variable.
- It **installs the selected runtime** using the PluginStore runtime interface from `sientia_model.model_repository.plugin_store`:
  - `PluginStore.install_runtime(runtime_name)`.
- It **uses PluginStore to obtain models from the store**, instead of constructing them from the local ML template / mlops library:
  - Pipelines call `PluginStore.get_model(...)` to obtain `SientiaModel` instances.
  - The previous “in-repo model implementation plus mlops library” path is removed.

Out of scope:

- Changing Temporal workflow semantics (queues, retry policies, etc.).
- Replacing the existing MLflow-based tracking and reporting; these remain the responsibility of `ModelRepository` and the reporting utilities.

---

## Existing State Overview (model-manager)

The `sientia-dataops-model-manager` application currently:

- **Worker orchestration** (`model_manager/worker/worker.py`)
  - Reads Temporal configuration from env (`TEMPORAL_HOST`, `TEMPORAL_NAMESPACE`, task queues).
  - Sets up observability (Prometheus, metrics, Sientia logger).
  - Builds connector configs for Postgres, MLflow, MinIO, MongoDB.
  - Instantiates `Activities` and schedules, then starts Temporal workers.
  - Does not validate or install any “runtime” concept before worker startup.

- **Training pipeline**
  - `Training` activity (`model_manager/activities/training.py`) coordinates:
    - Parameter validation (`validate_train_params`).
    - Training execution via `TrainingRepository`.
    - Saving trained models and artifacts via `ModelRepository` to MLflow.
  - `TrainingRepository` (`model_manager/utils/repository/training_repository.py`):
    - Loads and preprocesses CSV data (via `DataPreprocessor`).
    - Trains a local `LinearRegressionModel` defined in `model_manager.sientia.models`.
    - Computes metrics and builds a `TrainModelResult` for downstream steps.
  - `ModelRepository` (`model_manager/utils/repository/model_repository.py`):
    - Uses `ModelServing` and `Reports` to:
      - Generate reports and CSV artifacts.
      - Log parameters, metrics and models into MLflow.
      - Today, production models are identified using **stages** (for example `Production`) in the Model Registry; aliases like `@production` are not yet used.
      - MLflow-related logic here overlaps conceptually with the MLflow interactions implemented in `sientia-dataops-laborious_temporal`, which motivates extracting a **shared MLflow repository** into `sientia-dataops-library` (see `mlflow-shared-repository-migration-plan`).

- **Coupling to models and runtimes**
  - Training is tightly coupled to internal classes (`DataPreprocessor`, `LinearRegressionModel`).
  - No runtime installation; env assumed ready.
  - The system does not yet use PluginStore’s `get_model` or `install_runtime` capabilities.

**MLflow:** All MLflow operations (lookup, promotion, etc.) → [[mlflow-shared-repository-migration-plan|shared repository]]. Model Manager uses the interface; it does not implement these concepts.

### Current vs Target — High-level Flow

```mermaid
flowchart LR
    subgraph currentState [Current State — Model Manager]
        direction TB
        WorkerMM["Temporal Worker\n(model_manager/worker.py)"]
        ActivitiesMM["Activities\n(training, cleanup, etc.)"]
        TrainRepo["TrainingRepository\n(local DataPreprocessor + LinearRegressionModel)"]
        ModelRepoMM["ModelRepository\n(MLflow + reports)"]

        WorkerMM -->|"Temporal tasks"| ActivitiesMM
        ActivitiesMM -->|"train_model activity"| TrainRepo
        TrainRepo -->|"TrainModelResult"| ModelRepoMM
        ModelRepoMM -->|"experiments, runs, artifacts"| MLflowMM["MLflow Server"]
    end

    subgraph targetState [Target State — Model Manager]
        direction TB
        WorkerMM2["Temporal Worker\n+ Runtime bootstrap\n(read RUNTIME, install runtime)"]
        ActivitiesMM2["Activities\n(+ ModelProvider)"]
        PluginStoreNode["PluginStore\n(Gitea model store)"]
        SientiaWrapper["SientiaModel wrapper\n(from store.get_model)"]
        ModelRepoMM2["ModelRepository\n(MLflow + reports)"]

        WorkerMM2 -->|"install_runtime(RUNTIME)"| PluginStoreNode
        WorkerMM2 -->|"Temporal tasks"| ActivitiesMM2
        ActivitiesMM2 -->|"get_training_model"| PluginStoreNode
        PluginStoreNode -->|"SientiaModel instance"| SientiaWrapper
        ActivitiesMM2 -->|"call train(...) on wrapper"| SientiaWrapper
        SientiaWrapper -->|"TrainModelResult-compatible data"| ModelRepoMM2
        ModelRepoMM2 -->|"experiments, runs, artifacts"| MLflowMM2["MLflow Server"]
    end
```

### Current vs Target — Training Hot Path (Code Sketch)

Current hot path in `TrainingRepository.train` (simplified):

```python
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
data = _ensure_date_column_parsed(data, params)
data = self._configure_datetime_index(data, params)

process_data = self._init_data_preprocessor(params)
process_data.fit(data)
data_view = process_data.transform(data)

x_train, x_test, y_train, y_test = split_train_test(...)

regr = LinearRegressionModel(
    target_variable=params.target_variable,
    variable_columns=params.variable_columns,
    degree=params.degree,
    interaction_only=params.interaction_only,
)
regr.fit(pd.concat([x_train, y_train], axis=1))
```

Target hot path with PluginStore + `SientiaModel` (conceptual):

```python
data = load_and_preprocess(uploaded_file, params)          # keep existing preprocessing
train_df, val_df = build_train_val_splits(data, params)   # explicit train/val

wrapper = model_provider.get_training_model(
    model_name=params.model_name,
    runtime=os.environ["RUNTIME"],
    opt_params={"env": params.environment},
    model_kwargs={},
    data_model_kwargs={},
)

wrapper.train(
    train_data=train_df,
    val_data=val_df,
    target=params.target_variable,
)

# From here, metrics and artifacts are computed based on wrapper outputs
```

---

## Requirements Mapping

### Functional Requirements

| ID    | Requirement                                  | Description                                                                                          | Impacted Areas                                                                 |
|-------|----------------------------------------------|------------------------------------------------------------------------------------------------------|--------------------------------------------------------------------------------|
| FR-01 | Runtime detection via env var                | Before starting workers, read `RUNTIME` env var and fail fast if missing or empty                   | `model_manager/worker/worker.py`                                               |
| FR-02 | Runtime installation via PluginStore         | Install the runtime defined in `RUNTIME` using `PluginStore.install_runtime(runtime_name)`          | `worker.py`, `sientia_model.model_repository.plugin_store.PluginStore`        |
| FR-03 | PluginStore-based model retrieval            | Use `PluginStore.get_model(...)` to obtain `SientiaModel` instances for training                    | `TrainingRepository` (or new adapter), `Training` activity                     |
| FR-04 | Remove mlops-library-based construction      | Do not construct models directly from `model_manager.sientia.models`; remove the mlops path  | `TrainingRepository`, `model_manager.sientia.models`, dependency graph        |
| FR-05 | Use shared MLflow repository                | Delegate all MLflow ops (runs, metrics, artifacts, production lookup) to `SientiaMLflowRepository`; keep reports and `TrainModelResult` in Model Manager | `ModelRepository`, [[mlflow-shared-repository-migration-plan]] |
| FR-06 | Centralized PluginStore configuration       | Configure PluginStore (Gitea URL, repo, auth, PyPI mirror) via env/config | Config helper, `worker.py`                                    |

### Non-Functional Requirements

| ID     | Requirement                              | Description                                                                                                       | Impacted Areas                                             |
|--------|------------------------------------------|-------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------|
| NFR-01 | Robust runtime bootstrap                 | If runtime installation fails, the worker must not start; errors must be explicit in logs/metrics                 | `worker.py`, PluginStore error handling                    |
| NFR-02 | Single path                          | Use PluginStore + SientiaModel path only; no toggling or parallel old path | `worker.py`, Training activities                         |
| NFR-03 | Observability and debuggability          | Logs and metrics must cover runtime detection, install attempts and PluginStore interactions                      | Logging around runtime and PluginStore                     |
| NFR-04 | Testability                              | Unit and integration tests must cover runtime detection, install, and model retrieval                             | `tests/worker`, `tests/activities`, new PluginStore tests |
| NFR-05 | Security                                 | No credentials hard-coded in source; rely on env/secret management                                                | Config helpers, deployment manifests                       |

---

## Target Architecture

### Runtime bootstrap

- **New required env var:** `RUNTIME` (name of the runtime; must match a runtime in the store index).
- **New PluginStore configuration env vars** (names to be finalized):
  - `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`
  - Optional: `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `PYPI_INDEX_URL`, `PYPI_USERNAME`, `PYPI_PASSWORD`
- **Worker startup sequence** in `main()`:
  1. Initialize logger and basic metadata as today.
  2. Read `RUNTIME` from env; if missing/empty → log critical error and exit.
  3. Build `PluginStore` instance using configuration env vars.
  4. Call `store.install_runtime(runtime_name=RUNTIME)`; on failure → log error, set `APP_UP` metric to 0 and exit.
  5. Only after successful runtime installation → build connector configs, instantiate Activities, connect to Temporal.

### PluginStore usage in training

- **ModelProvider abstraction** (e.g. `model_manager/utils/model_provider.py`):
  - Holds a `PluginStore` instance and logger.
  - Exposes `get_training_model(model_name, runtime, opt_params, model_kwargs, data_model_kwargs) -> SientiaModel` via `store.get_model(...)`.
- **TrainingRepository integration:**
  - Prepare input data using existing preprocessing logic.
  - Ask ModelProvider for a `SientiaModel` instance (`model_name` from params, `runtime` from `RUNTIME`).
  - Call the public `train(...)` method of the wrapper; collect outputs and metadata only through the public API.
  - Build `TrainModelResult` from train/test splits, predictions, and artifacts required by `ModelRepository`.
  - Model behavior lives in `SientiaModel`; manager focuses on orchestration and reports.

---

## Implementation Plan

### Phase 0 — Design Alignment

- **P0-01**: Confirm with stakeholders:
  - The expected values and semantics for `RUNTIME` (naming convention, mapping to store runtimes).
  - Whether the manager will ever need to support more than one runtime per process (current assumption: no).
- **P0-02**: Decide the config strategy:
  - Pure environment variables vs a config file + env overrides.
- **P0-03**: Validate how `model_name` will be passed into the training workflow:
  - Confirm or extend `TrainModelParams` to carry `model_name` and runtime-related fields as needed.

### Phase 1 — Runtime Detection and Installation

- **P1-01**: Extend `worker.py` docs and configuration:
  - Add `RUNTIME` to the environment variables list in the module docstring.
  - Document failure behavior when `RUNTIME` is missing or empty.
- **P1-02**: Implement `build_plugin_store_from_env` helper:
  - New function that:
    - Reads `STORE_*` and `PYPI_*` env vars.
    - Creates and returns a `PluginStore` instance with appropriate logger.
  - Place it either in `worker.py` or in a dedicated utility module (e.g. `model_manager/utils/plugin_store_config.py`).
- **P1-03**: Integrate runtime installation into `main()`:
  - Before any Temporal client initialization:
    - Read `runtime_name = os.getenv("RUNTIME")`.
    - Build PluginStore via the helper.
    - Call `install_runtime(runtime_name)`.
  - Log:
    - Start and end of runtime installation.
    - List of installed requirements (name and version).
  - On failure:
    - Emit a clear message (including runtime name and store repo).
    - Mark the app as DOWN in metrics.
    - Exit with non-zero status.

### Phase 2 — ModelProvider and Training Integration

- **P2-01**: Introduce `ModelProvider` abstraction:
  - Implement `ModelProvider` with:
    - A reference to the shared `PluginStore`.
    - Methods for retrieving models for training:
      - `get_training_model(model_name, runtime, opt_params, model_kwargs, data_model_kwargs)`.
  - Ensure it logs:
    - Model and runtime names.
    - Cache hits/misses when appropriate.
- **P2-02**: Wire ModelProvider into Activities:
  - Update the construction of `Activities` in `worker.py` to accept:
    - A `ModelProvider` or a `PluginStore` that can be wrapped inside the `Training` activity.
  - Update `Training.__init__` signature to accept this new dependency and store it as an attribute.
- **P2-03**: Adjust `TrainingRepository` to use `SientiaModel`:
  - In `train`:
    1. Keep or refine current CSV loading and preprocessing pipeline (DataPreprocessor, split into train/test).
    2. Use `ModelProvider` to retrieve the model:
       - `model_name` from `TrainModelParams`.
       - `runtime` from `RUNTIME`.
       - `opt_params` mirroring how `test_plugin_store.py` interacts with models.
    3. Call the model’s `train` method and produce the same core artifacts the current code expects:
       - Train/test splits.
       - Metrics and predictions for `TrainModelResult`.
  - In `after_train_calculation`:
    - Either re-use metrics from the SientiaModel, or continue calculating MSE/MAE/R² as a verification step.

### Phase 3 — Removing mlops Dependencies

- **P3-01**: Identify all usage points of:
  - `LinearRegressionModel`.
  - `DataPreprocessor` where behavior overlaps with what SientiaModel already does.
- **P3-02**: Remove mlops-based code:
  - Remove unused mlops-library-based training hooks.
  - Simplify `TrainingRepository` to delegate as much as possible to SientiaModel logic.
  - Use PluginStore + `SientiaModel` path only (no migration flag or parallel old path).

### Phase 4 — ModelRepository, Shared MLflow Repository and Reporting Alignment

- **P4-01**: Ensure `TrainModelResult` is correctly populated:
  - Confirm that:
    - `x_train`, `x_test`, `y_train`, `y_test`, `y_pred`, `y_train_pred` are provided by the new flow.
    - Any fields required by `_generate_report` and `_save_run` remain available.
- **P4-02**: Delegate all MLflow operations to `SientiaMLflowRepository` (see [[mlflow-shared-repository-migration-plan]]); keep in Model Manager only reporting orchestration and `TrainModelResult` construction.
- **P4-03**: Validate that MLflow reports remain consistent:
  - Run a side-by-side comparison between:
    - A run from the previous (mlops-based) pipeline.
    - A PluginStore-based run for the same dataset/experiment using the shared repository.
  - Compare:
    - Logged parameters.
    - Metrics.
    - Artifacts (reports, CSVs, equation JSON).

---

## Testing Strategy

- **T1 — Unit tests**
  - Worker: Test behavior when `RUNTIME` is missing or empty; test that `install_runtime` is called with the correct runtime name (mock PluginStore).
  - ModelProvider: Test that it calls `PluginStore.get_model` with the expected arguments.
  - TrainingRepository: Test that it uses the model returned by PluginStore instead of `LinearRegressionModel`.
- **T2 — Integration tests**
  - Set up a test store with a minimal model and runtime.
  - Run a full training workflow: verify runtime installation first; confirm model is retrieved and training completes; validate metrics and artifacts via MLflow.
- **T3 — Regression tests**
  - Run the same experiment once with the old pipeline and once with the PluginStore-based pipeline.
  - Compare key results (metrics and artifacts) to ensure differences are understood and acceptable.

---

## Rollout and Migration Strategy

- **R1 — Criar modelos com a nova arquitetura**
  - Garantir que o pipeline de modelagem (Factory/Warehouse/Store) consiga produzir modelos:
    - Encapsulados em wrappers que estendem `SientiaModel`.
    - Com interface estável para `train`, `retrain` (quando existir), `predict`/`transform` e `store_model`.
  - Publicar um conjunto inicial de modelos “pilot” no store que será consumido pelo Model Manager.

- **R2 — Subir um ou mais runtimes para esses modelos**
  - Configurar e instalar runtimes específicos para os novos modelos, alinhados ao `RUNTIME` esperado pelo Model Manager:
    - Verificar que cada runtime contém todas as dependências necessárias (via PluginStore / runtime installer).
    - Validar que, em um ambiente de teste, o runtime consegue:
      - Instalar bibliotecas.
      - Carregar o wrapper via PluginStore e executar pelo menos um ciclo de treino de ponta a ponta.

- **R3 — Colocar os novos modelos para rodar no Model Manager**
  - Integrar o uso de `PluginStore.get_model(...)` na pipeline de treino:
    - Direcionar um subconjunto de fluxos de treinamento para os modelos “pilot” oriundos do store.
    - Monitorar:
      - Estabilidade dos workers.
      - Tempo de treino e consumo de recursos.
      - Artefatos e métricas geradas no MLflow.
  - Quando estáveis, coordenar com as equipes de produto/negócio para considerar esses modelos como candidatos a produção e, quando já estiverem usando MLflow 3+ com wrappers, promover esses modelos para produção usando aliases (`@production`).

- **R4 — Migrar progressivamente os demais modelos**
  - Definir uma ordem de migração por domínio/família de modelo (por exemplo: modelos de predição de série temporal, modelos de classificação, etc.):
    - Para cada modelo legado:
      - Criar/ajustar o wrapper `SientiaModel` correspondente no Factory/Warehouse.
      - Garantir que o modelo passe a ser entregue pelo store e consumido via PluginStore.
      - Executar o ciclo de testes (T1–T3) descrito na seção anterior.
  - Após migrar todas as famílias de modelo:
    - Remover o caminho de código que instancia diretamente `LinearRegressionModel` e demais classes locais.
    - Limpar variáveis de configuração relacionadas à pipeline antiga (mlops library local).
    - Assumir como padrão único: Model Manager treinando apenas modelos vindos do store, via wrappers `SientiaModel` e, quando aplicável, com resolução de produção por aliases em MLflow 3+.

---

## Potential Model Library Changes

- **Clarifying the SientiaModel training interface:**
  - Ensure that `SientiaModel` provides a stable way to train (including validation split handling).
  - A clear contract for returning predictions and metrics.
  - Access to any internal state needed for `TrainModelResult` construction.
- **Improving PluginStore ergonomics:**
  - Optionally add a helper to build `PluginStore` from environment variables (reusable across applications).
  - More structured exceptions for missing runtime definitions, missing models, network and authentication issues.
- These changes should be coordinated so that model manager and any other consumers can rely on a consistent, documented behavior.

---

## Related Documents

- [[mlflow-shared-repository-migration-plan|MLflow Shared Repository Migration Plan]] — Concepts implemented in the common interface (production lookup, wrapper loading, promotion, etc.)
- [[analytics-implementation-plan|Runtime Analytics Helm Implementation Plan]]
- [[analytics|Runtime Analytics Architecture and Analysis]]
- [[../model-plugin-system/06-end-to-end-flow|Model Plugin System — End-to-End Flow]]