--- 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]]