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

22 KiB
Raw Permalink Blame History

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 PluginStores 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 models 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 (T1T3) 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.