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

tags, created, modified, created_by, modified_by, status

tags	created	modified	created_by	modified_by	status
engineering sientia runtime-system model-manager plugin-store migration-plan	2026-03-02	2026-03-02	Vitor Pimentel	Vitor Pimentel	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 — What this migration must achieve
2. #Existing State Overview (model-manager) — Current responsibilities and coupling points
3. #Requirements Mapping — Functional and non-functional requirements
4. #Target Architecture — Desired runtime and model-loading architecture
5. #Implementation Plan — Phased, detailed changes to apply
6. #Testing Strategy — How to validate the new behavior
7. #Rollout and Migration Strategy — How to safely roll out and deprecate old paths
8. #Potential Model Library Changes — Expected impact on sientia-model-library
9. #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. Model Manager uses the interface; it does not implement these concepts.

Current vs Target — High-level Flow

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):

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):

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 — Concepts implemented in the common interface (production lookup, wrapper loading, promotion, etc.)
• analytics-implementation-plan
• analytics
• ../model-plugin-system/06-end-to-end-flow