vitor-aignosi
1ce8b9d3a7
Enhance environment configuration and update dependencies - Added new environment variables for PluginStore and MLflow configuration in `.env.example`, including `RUNTIME`, `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`, `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `STORE_CACHE_TTL_SECONDS`, `PYPI_SERVER`, `PYPI_USERNAME`, and `PYPI_PASSWORD`. - Updated `git-requirements-mapping.txt` to reflect changes in repository names. - Modified `requirements-light.txt` and `requirements.txt` to upgrade `sientia-dataops-library` to version 1.12.0 and `sientia-mlops-library` to version 0.8.1. - Updated `values.yaml` to include new environment variables for worker runtime and PluginStore configuration. - Refactored E2E tests to utilize new MLflow repository stubs and PluginStore mocks for improved testing accuracy.
306 lines
17 KiB
Markdown
306 lines
17 KiB
Markdown
---
|
||
tags:
|
||
- engineering
|
||
- sientia
|
||
- runtime-system
|
||
- laborious-temporal
|
||
- plugin-store
|
||
- migration-plan
|
||
created: 2026-03-02
|
||
modified: 2026-03-02
|
||
created_by: Vitor Pimentel
|
||
modified_by: Vitor Pimentel
|
||
status: draft
|
||
---
|
||
|
||
# Sientia Laborious Temporal — PluginStore & Wrapper Migration Plan
|
||
|
||
> Migration plan for evolving `sientia-dataops-laborious_temporal` from direct MLflow model loading to a runtime-aware architecture that uses Sientia model wrappers (`SientiaModel`) via their public methods, aligned with the runtime strategy.
|
||
|
||
## Summary
|
||
|
||
1. [[#Objectives and Scope|Objectives and Scope]] — What this migration must achieve
|
||
2. [[#Existing State Overview (laborious_temporal)|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 interaction 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 the changes
|
||
8. [[#Related Documents|Related Documents]] — Cross-links to supporting documents
|
||
|
||
---
|
||
|
||
## Objectives and Scope
|
||
|
||
This migration focuses on the `sientia-dataops-laborious_temporal` application and aims to:
|
||
|
||
- Keep the **runtime-aware deployment model** consistent with the rest of the runtime system (Helm + `RUNTIME` env var, runtime installation via PluginStore).
|
||
- Ensure that **all interactions with models use the public methods of the Sientia wrapper** (`SientiaModel`):
|
||
- Use `SientiaModel.train(...)` and `retrain(...)` for training and retraining flows.
|
||
- Use `SientiaModel.predict(...)` and `SientiaModel.transform(...)` for inference and preprocessing.
|
||
- **Use the shared MLflow repository** (`SientiaMLflowRepository`) for all MLflow operations (load, runs, artifacts, promotion, production lookup, metadata logging); do not implement these in Laborious.
|
||
|
||
Out of scope:
|
||
|
||
- Replacing MLflow as the tracking and registry backend.
|
||
- Redesigning Temporal workflows (queues, retry policies) beyond what is required for the new model interaction style.
|
||
|
||
---
|
||
|
||
## Existing State Overview (laborious_temporal)
|
||
|
||
Key components in `sientia-dataops-laborious_temporal`:
|
||
|
||
- **MLflow activities** (`laborious/activities/mlflow.py`)
|
||
- `MLFlow` class exposes Temporal activities for:
|
||
- `request_transform` — loads transformation models from MLflow and applies them to input data.
|
||
- `request_predict` — loads predictive models from MLflow and generates predictions.
|
||
- `retrain_model` — orchestrates retraining using historical data stored in MinIO and MLflow registry.
|
||
- `update_production_model` — promotes new versions to production.
|
||
- `get_reference_data` — fetches evaluation/reference datasets from model artifacts.
|
||
- These activities delegate ML-specific work to `MLFlowRepository`.
|
||
|
||
- **MLflow repository** (`laborious/utils/repository/model_repository.py`)
|
||
- `MLFlowRepository` encapsulates the interaction with MLflow:
|
||
- Model discovery and run resolution (`get_model_run_id`, `get_model_uri`, `get_experiment`, etc.).
|
||
- Artifact download and loading for both transformer and prediction models.
|
||
- Model caching and retention (`get_model`, `get_cached_operation`).
|
||
- Transformation and prediction entry points:
|
||
- `transform(...)` wraps `get_cached_operation(..., operation='transform')`.
|
||
- `predict(...)` wraps `get_cached_operation(..., operation='predict')`.
|
||
- Retraining orchestration (`fit_models`, `create_new_experiment`, `retrain_model`, `update_production_model`).
|
||
- Today:
|
||
- Models are loaded via MLflow flavors: sklearn, pyfunc, pytorch.
|
||
- When `flavor == 'pyfunc'` and `load_wrapper=True`, the repository loads a wrapper via:
|
||
- `raw_model = mlflow.pyfunc.load_model(artifact_path)`
|
||
- `model = raw_model._model_impl.python_model`
|
||
- Production models are resolved using **stages** in the Model Registry (for example, selecting the latest version in stage `Production`); **aliases such as `@production` are not used yet**, and models are registered explicitly as part of the current retrain/promotion flows.
|
||
- The wrapper’s `_model_impl` class does not extend `SientiaModel`.
|
||
- All this MLflow-specific logic is local to `laborious_temporal` and partially duplicated in `sientia-dataops-model-manager`, which motivates the extraction of a shared MLflow repository in `sientia-dataops-library` (see `mlflow-shared-repository-migration-plan`).
|
||
|
||
**MLflow:** All MLflow ops → [[mlflow-shared-repository-migration-plan|shared repository]]. Laborious uses the interface; `SientiaModel` lifecycle is in sientia-model-library.
|
||
|
||
### Current vs Target — High-level Flow
|
||
|
||
```mermaid
|
||
flowchart LR
|
||
subgraph current [Current State — Laborious Temporal]
|
||
direction TB
|
||
TemporalWorker["Temporal Worker"]
|
||
MlflowActivities["MLFlow Activities\nrequest_transform / request_predict / retrain_model"]
|
||
MLFlowRepositoryNode["MLFlowRepository"]
|
||
MLflowRegistry["MLflow Tracking + Registry"]
|
||
RawModel["Loaded Model\n(sklearn / pyfunc / pytorch)"]
|
||
|
||
TemporalWorker -->|"start workflow\n(Temporal)"| MlflowActivities
|
||
MlflowActivities -->|"call transform()/predict()/retrain_model()"| MLFlowRepositoryNode
|
||
MLFlowRepositoryNode -->|"search_model_versions()\ncurrent_stage == 'Production'"| MLflowRegistry
|
||
MLFlowRepositoryNode -->|"mlflow.*.load_model(model_uri)"| RawModel
|
||
MLFlowRepositoryNode -->|"raw_model.predict(data)\nraw_model.fit(data)"| RawModel
|
||
end
|
||
|
||
subgraph target [Target State — Laborious Temporal]
|
||
direction TB
|
||
TemporalWorker2["Temporal Worker"]
|
||
MlflowActivities2["MLFlow Activities\n(same APIs)"]
|
||
MLFlowRepositoryNode2["MLFlowRepository\n(wrapper-aware)"]
|
||
MLflowRegistry2["MLflow Tracking + Registry\n(aliases enabled)"]
|
||
SientiaWrapperNode["Wrapper Instance\n(extends SientiaModel)"]
|
||
|
||
TemporalWorker2 -->|"start workflow\n(Temporal)"| MlflowActivities2
|
||
MlflowActivities2 -->|"call transform()/predict()/retrain_model()"| MLFlowRepositoryNode2
|
||
MLFlowRepositoryNode2 -->|"get_model_version_by_alias('production')\n& models:/name@production"| MLflowRegistry2
|
||
MLFlowRepositoryNode2 -->|"mlflow.pyfunc.load_model(...)"| SientiaWrapperNode
|
||
MLFlowRepositoryNode2 -->|"wrapper.transform(...)\nwrapper.predict(...)\nwrapper.train()/retrain()"| SientiaWrapperNode
|
||
SientiaWrapperNode -->|"store_model(...)\n(auto-register + update alias)"| MLflowRegistry2
|
||
end
|
||
```
|
||
|
||
### Current vs Target — Retrain Hot Path (Code Sketch)
|
||
|
||
Current retrain flow inside `MLFlowRepository.fit_models` / `retrain_model` (simplified):
|
||
|
||
```python
|
||
data_model, _ = await self.download_model(
|
||
model_name=model_name,
|
||
metadata=metadata,
|
||
model_type="transform",
|
||
flavor=transform_flavor,
|
||
load_wrapper=(transform_flavor == "pyfunc"),
|
||
)
|
||
|
||
prediction_model, _ = await self.download_model(
|
||
model_name=model_name,
|
||
metadata=metadata,
|
||
model_type="predict",
|
||
flavor=predict_flavor,
|
||
load_wrapper=(predict_flavor == "pyfunc"),
|
||
)
|
||
|
||
treated_data_candidate = data_model.fit(data) # or data_model.predict(data)
|
||
...
|
||
prediction_model.fit(retrain_dataset) # direct fit on underlying model
|
||
```
|
||
|
||
Target retrain flow when using wrappers that extend `SientiaModel`:
|
||
|
||
```python
|
||
wrapper, _ = await self.download_model(
|
||
model_name=model_name,
|
||
metadata=metadata,
|
||
# model_type="predict", # wrapper owns both transformer + model stack
|
||
# flavor="pyfunc", flavor will always be pyfunc, this parameter will be removed
|
||
# load_wrapper=True, wrapper will always be loaded from _model_impl
|
||
)
|
||
|
||
# First-time training or full retrain using public API
|
||
wrapper.train(
|
||
train_data=train_df, # features + target
|
||
val_data=val_df,
|
||
target=target_name,
|
||
)
|
||
|
||
# Incremental retrain (when appropriate)
|
||
wrapper.retrain(full_retrain_df)
|
||
|
||
transformed_df, trans_meta = wrapper.transform(raw_df)
|
||
pred_df, pred_meta = wrapper.predict({}, transformed_df, params={})
|
||
```
|
||
|
||
---
|
||
|
||
## Requirements Mapping
|
||
|
||
### Functional Requirements
|
||
|
||
| ID | Requirement | Description | Impacted Areas |
|
||
| ----- | -------------------------------------------------- | ------------------------------------------------------------------------------------------------ | ------------------------------------------------------ |
|
||
| FR-01 | Runtime detection and installation | Align worker startup with `RUNTIME` and runtime installation via PluginStore | Worker |
|
||
| FR-02 | Wrapper‑based training using public API | Retraining must call the public `retrain(...)` method of `SientiaModel` | `fit_models`, `retrain_model` paths |
|
||
| FR-03 | Wrapper‑based inference using public API | Inference must call `predict(...)` and `transform(...)` on the wrapper; obtain wrappers via `SientiaMLflowRepository` | Activities, shared repository |
|
||
| FR-04 | Use shared MLflow repository | Delegate all MLflow operations (load, runs, artifacts, promotion, production lookup) to `SientiaMLflowRepository`; do not implement in Laborious | [[mlflow-shared-repository-migration-plan]] |
|
||
| FR-05 | Model configuration for training/retraining | `model_config` must define `target` and `retention_minutes`; all models are pyfunc + wrapper (no flavor selection) | `model_config` structures |
|
||
|
||
### Non-Functional Requirements
|
||
|
||
| ID | Requirement | Description | Impacted Areas |
|
||
| ------ | -------------------------------- | ------------------------------------------------------------------------------------------------------------------------ | --------------------------------------- |
|
||
| NFR-01 | Consistent public API usage | All wrapper interactions must go through public `SientiaModel` methods; metadata logging is handled by the shared repository | Activities |
|
||
| NFR-02 | Fail fast when wrappers unavailable | Where wrappers are not yet available, raise an explicit error requesting model update | Shared repository, config |
|
||
| NFR-03 | Testability | Enable unit tests to validate wrapper-based flows and integration with shared repository | `tests/laborious` |
|
||
| NFR-04 | Operational safety | Retraining and promotion behavior must remain auditable and robust | Retrain & promotion flows |
|
||
|
||
---
|
||
|
||
## Target Architecture
|
||
|
||
### Wrapper-centric model interactions
|
||
|
||
The target state for `laborious_temporal` is:
|
||
|
||
- All training and retraining logic goes through the wrapper’s public methods: `train(...)` and `retrain(...)`.
|
||
- Inference and transformation use `transform(df)` and `predict(context, df, params)`.
|
||
- All MLflow operations (load, runs, artifacts, promotion, production lookup) go through `SientiaMLflowRepository` (see [[mlflow-shared-repository-migration-plan]]).
|
||
|
||
### Use of Shared MLflow Repository
|
||
|
||
Laborious obtains wrappers and performs all MLflow operations via `SientiaMLflowRepository`. The shared repository (see [[mlflow-shared-repository-migration-plan]]) owns: alias-based URIs, pyfunc loading, wrapper extraction, promotion, runs, artifacts, and metadata logging. Laborious activities call `repo.load_wrapper(...)` and then the wrapper’s public methods (`transform`, `predict`, `train`, `retrain`); they do not implement MLflow logic.
|
||
|
||
---
|
||
|
||
## Implementation Plan
|
||
|
||
### Phase 0 — Design and Configuration Alignment
|
||
|
||
- **P0-01**: Catalog model types and flavors used by `laborious_temporal`:
|
||
- For each active model:
|
||
- Flavor will be always pyfunc
|
||
- Whether a `_model_impl` wrapper is already present and extends `SientiaModel`.
|
||
- **P0-02**: Define configuration fields in `model_config` for wrapper usage:
|
||
- Example:
|
||
- `target` field for training/retraining (already partially present).
|
||
- `retention_minutes` field for model retention in minutes (unchanged).
|
||
|
||
### Phase 1 — Inference via SientiaModel Public API (using shared repository)
|
||
|
||
- **P1-01**: Replace `download_model` with `SientiaMLflowRepository.load_wrapper(...)`; remove local MLflow loading logic.
|
||
- **P1-02**: Update `get_cached_operation` to call `wrapper.transform(data)` and `wrapper.predict({}, data)`; unpack returned metadata; metadata logging is handled by the shared repository.
|
||
- **P1-03**: Ensure activities (`request_transform`, `request_predict`) remain unchanged externally (inputs/outputs unchanged).
|
||
|
||
### Phase 2 — Retraining via SientiaModel.retrain
|
||
|
||
- **P2-01**: Refactor `fit_models` to call `wrapper.retrain(data)` and `wrapper.train(...)`; use `SientiaMLflowRepository` for runs, metrics, artifacts, and promotion (no local MLflow logic).
|
||
|
||
### Phase 3 — MLflow Logging and Promotion (via shared repository)
|
||
|
||
- **P3-01**: In `create_new_experiment`, use the wrapper’s `store_model(...)` for model artifacts; use `SientiaMLflowRepository` for runs, metrics, and any additional MLflow operations.
|
||
- **P3-02**: Use `SientiaMLflowRepository.promote_to_alias(...)` for production promotion; model registration is auto-handled by wrappers.
|
||
|
||
### Phase 4 — Runtime Alignment
|
||
|
||
- **P4-01**:`laborious_temporal` is also deployed via the runtime-aware Helm chart:
|
||
- Read `RUNTIME` env var.
|
||
- Install runtime via PluginStore before starting Temporal workers.
|
||
- **P4-02**: Standardize worker queue name as `{project_name}-{runtime}-queue`.
|
||
- **P4-03**: Fix quality pipelines to a single runtime (to be decided).
|
||
|
||
---
|
||
|
||
## Testing Strategy
|
||
|
||
- **T1 — Unit tests**
|
||
- Add tests in `tests/laborious/utils/repository/test_model_repository.py` to cover:
|
||
- Wrapper-based `get_cached_operation` for both `transform` and `predict` (using shared repository).
|
||
- Wrapper-based `fit_models` and `retrain_model` paths calling `train` and `retrain` respectively.
|
||
- Code coverage must be 100%.
|
||
- **T2 — Integration tests**
|
||
- Use existing end-to-end tests under `e2e/`:
|
||
- Configure a model with `SientiaModel` wrapper (loaded via shared repository).
|
||
- Run full prediction and retrain workflows; compare predictions, retrain outcomes, and MLflow artifacts.
|
||
|
||
---
|
||
|
||
## Rollout and Migration Strategy
|
||
|
||
- **R1 — Create models with the new architecture**
|
||
- Update or create models in the modeling pipeline so they:
|
||
- Use wrappers that extend `SientiaModel`.
|
||
- Correctly implement `train`, `retrain`, `predict`, `transform`, and `store_model`.
|
||
- Produce structured metadata in `transform_meta` and `pred_meta`.
|
||
- Publish these models to a test store (or dedicated branch/experiment) for initial validation.
|
||
|
||
- **R2 — Provision runtimes for the new models**
|
||
- Configure and install dedicated runtimes for the new models:
|
||
- Ensure runtime dependencies (Python and system libraries) are available via PluginStore/runtime installer.
|
||
- Validate that each runtime can:
|
||
- Load the wrapper through MLflow.
|
||
- Execute `transform` and `predict` end-to-end on sample data.
|
||
|
||
- **R3 — Run new models in real workflows**
|
||
- Integrate the new wrapper-based models into real Laborious workflows, initially in non-critical environments:
|
||
- Route only a subset of flows or entities to the new models.
|
||
- Monitor logs (including metadata), business metrics, and retraining/promotion behavior.
|
||
- Promote these models to production using aliases (`@production`) in MLflow 3+.
|
||
|
||
- **R4 — Migrate remaining models progressively**
|
||
- Define migration waves by model family:
|
||
- For each existing model:
|
||
- Create or adapt a `SientiaModel` wrapper.
|
||
- Provision the corresponding runtime.
|
||
- Execute the test cycle (T1–T2) from the previous section.
|
||
- Update aliases so production traffic uses the new wrapper.
|
||
- After all models are migrated:
|
||
- Remove legacy stage-based paths (`Production`) and non-wrapper models.
|
||
- Simplify the codebase to assume wrappers + aliases only.
|
||
|
||
---
|
||
|
||
## Related Documents
|
||
|
||
- [[mlflow-shared-repository-migration-plan|MLflow Shared Repository Migration Plan]] — Concepts implemented in the common interface (production lookup, wrapper loading, promotion, metadata logging, etc.)
|
||
- [[model-manager-plugin-store-migration-plan|Model Manager PluginStore Migration Plan]]
|
||
- [[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]]
|
||
|