feat: integrate PluginStore and MinIO repository into model manager activities
- Added PluginStore integration for model management. - Replaced StorageRepository with MinIORepository in Activities, Cleanup, and Training classes. - Updated training logic to handle validation files and improved data management. - Enhanced configuration for MinIO and PluginStore in connectors. - Removed deprecated model repository and storage repository files. - Updated environment variable handling for new configurations.
This commit is contained in:
vitor-aignosi
385
model-manager-plugin-store-migration-plan.md
Normal file
385
model-manager-plugin-store-migration-plan.md
Normal file
@@ -0,0 +1,385 @@
|
||||
---
|
||||
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]]
|
||||
|
||||
@@ -6,12 +6,14 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from model_manager.activities.cleanup import Cleanup
|
||||
from model_manager.activities.experiment_tracking import ExperimentTracking
|
||||
from model_manager.activities.training import Training
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
|
||||
class Activities(ExperimentTracking, Training, Cleanup):
|
||||
@@ -40,8 +42,10 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
postgres_config: dict[str, Any],
|
||||
mlflow_config: dict[str, Any],
|
||||
minio_config: dict[str, Any],
|
||||
plugin_store: PluginStore,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
):
|
||||
"""
|
||||
Initialize the Activities orchestrator with all required configurations.
|
||||
@@ -62,9 +66,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
Raises:
|
||||
Exception: If any parent class initialization fails
|
||||
"""
|
||||
metrics_controller = MetricsController(
|
||||
logger=logger,
|
||||
)
|
||||
|
||||
|
||||
ExperimentTracking.__init__(
|
||||
self,
|
||||
@@ -80,30 +82,40 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
self.model_repository = ModelRepository(
|
||||
url=mlflow_config['url'],
|
||||
self.mlflow_repository = SientiaMLflowRepository(
|
||||
host=mlflow_config['url'],
|
||||
username=mlflow_config['username'],
|
||||
password=mlflow_config['password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
self.storage_repository = StorageRepository(
|
||||
endpoint_url=minio_config['endpoint_url'],
|
||||
# MinIO repository used for all object storage operations
|
||||
endpoint_url = minio_config['endpoint_url']
|
||||
# MinioRepository expects the endpoint without scheme
|
||||
if endpoint_url.startswith('http://'):
|
||||
endpoint = endpoint_url.removeprefix('http://')
|
||||
elif endpoint_url.startswith('https://'):
|
||||
endpoint = endpoint_url.removeprefix('https://')
|
||||
else:
|
||||
endpoint = endpoint_url
|
||||
|
||||
self.minio_repository = MinioRepository(
|
||||
endpoint=endpoint,
|
||||
access_key=minio_config['access_key'],
|
||||
secret_key=minio_config['secret_key'],
|
||||
region=minio_config['region'],
|
||||
use_ssl=minio_config['use_ssl'],
|
||||
max_retry_attempts=minio_config['max_retry_attempts'],
|
||||
retry_mode=minio_config['retry_mode'],
|
||||
connect_timeout=minio_config['connect_timeout'],
|
||||
read_timeout=minio_config['read_timeout'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
secure=minio_config['use_ssl'],
|
||||
)
|
||||
|
||||
Training.__init__(
|
||||
self,
|
||||
model_repository=self.model_repository,
|
||||
storage_repository=self.storage_repository,
|
||||
mlflow_repository=self.mlflow_repository,
|
||||
plugin_store=plugin_store,
|
||||
minio_repository=self.minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
@@ -111,7 +123,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
|
||||
Cleanup.__init__(
|
||||
self,
|
||||
storage_repository=self.storage_repository,
|
||||
minio_repository=self.minio_repository,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
@@ -137,7 +149,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
# Logging here could cause issues if logger is already destroyed
|
||||
pass
|
||||
|
||||
async def shutdown(self):
|
||||
def shutdown(self):
|
||||
"""
|
||||
Gracefully shutdown all activities and clean up resources.
|
||||
|
||||
@@ -152,4 +164,4 @@ class Activities(ExperimentTracking, Training, Cleanup):
|
||||
"""
|
||||
ExperimentTracking.close(self)
|
||||
self.info('Postgres client closed')
|
||||
self.storage_repository.close()
|
||||
SientiaMonitoring.shutdown(self)
|
||||
|
||||
@@ -21,9 +21,9 @@ with workflow.unsafe.imports_passed_through():
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
|
||||
from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24'))
|
||||
DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true'
|
||||
@@ -41,7 +41,7 @@ class Cleanup(SientiaMonitoring):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
storage_repository: StorageRepository,
|
||||
minio_repository: MinioRepository,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
@@ -50,13 +50,13 @@ class Cleanup(SientiaMonitoring):
|
||||
Initialize Cleanup activity.
|
||||
|
||||
Args:
|
||||
storage_repository: Repository for MinIO operations
|
||||
minio_repository: Repository for MinIO operations
|
||||
logger: Logger instance for observability
|
||||
notification_handler: Handler for sending notifications
|
||||
metrics_controller: Controller for metrics emission
|
||||
"""
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
self.storage_repository = storage_repository
|
||||
self.minio_repository = minio_repository
|
||||
|
||||
# Configuration from environment variables
|
||||
self.retention_hours = RETENTION_HOURS
|
||||
@@ -113,15 +113,21 @@ class Cleanup(SientiaMonitoring):
|
||||
files_deleted = 0
|
||||
errors = []
|
||||
|
||||
# List objects in the specified bucket
|
||||
max_keys = self.max_keys_cleanup # Use environment variable for page size
|
||||
objects = self.storage_repository.list_bucket_objects(bucket_name, max_keys)
|
||||
# List objects in the specified bucket. MinioRepository applies BASE_PREFIX
|
||||
# internally; we request all objects under that prefix for this bucket.
|
||||
objects = await self.minio_repository.list_objects(
|
||||
prefix='',
|
||||
bucket=bucket_name,
|
||||
recursive=True,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
for obj_key in objects:
|
||||
files_scanned += 1
|
||||
|
||||
# Extract timestamp from filename
|
||||
match = self.minio_timestamp_pattern.match(obj_key)
|
||||
# Extract timestamp from the filename portion of the object key
|
||||
filename = obj_key.split('/')[-1]
|
||||
match = self.minio_timestamp_pattern.match(filename)
|
||||
if not match:
|
||||
self.debug(f'Skipping file without timestamp pattern: {obj_key}', metadata)
|
||||
continue
|
||||
@@ -137,10 +143,14 @@ class Cleanup(SientiaMonitoring):
|
||||
files_deleted += 1
|
||||
else:
|
||||
try:
|
||||
self.storage_repository.delete_file(bucket_name, obj_key)
|
||||
await self.minio_repository.delete_file(
|
||||
object_name=obj_key,
|
||||
bucket=None,
|
||||
metadata=metadata,
|
||||
)
|
||||
self.info(f'Deleted stale file: {obj_key}', metadata)
|
||||
files_deleted += 1
|
||||
except OSError as e:
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Failed to delete {obj_key}: {str(e)}'
|
||||
errors.append(error_msg)
|
||||
self.error(error_msg, metadata)
|
||||
|
||||
@@ -12,18 +12,20 @@ with workflow.unsafe.imports_passed_through():
|
||||
import traceback
|
||||
from typing import Any
|
||||
|
||||
import pandas as pd
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.repository.minio_repository import MinioRepository
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
|
||||
from model_manager.utils.exceptions import ModelTrainingError
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
from model_manager.utils.repository.training_repository import TrainingRepository
|
||||
from model_manager.utils.repository.data_manager_repository import DataManagerRepository
|
||||
|
||||
|
||||
class Training(SientiaMonitoring):
|
||||
@@ -38,8 +40,9 @@ class Training(SientiaMonitoring):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_repository: ModelRepository,
|
||||
storage_repository: StorageRepository,
|
||||
mlflow_repository: SientiaMLflowRepository,
|
||||
plugin_store: PluginStore,
|
||||
minio_repository: MinioRepository,
|
||||
logger: Logger,
|
||||
notification_handler: NotificationHandler,
|
||||
metrics_controller: MetricsController,
|
||||
@@ -52,9 +55,10 @@ class Training(SientiaMonitoring):
|
||||
notification_handler: Handler for sending notifications
|
||||
"""
|
||||
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
|
||||
self.training_repository = TrainingRepository(logger)
|
||||
self.model_repository = model_repository
|
||||
self.storage_repository = storage_repository
|
||||
self.data_manager_repository = DataManagerRepository(logger)
|
||||
self.mlflow_repository = mlflow_repository
|
||||
self.plugin_store = plugin_store
|
||||
self.minio_repository = minio_repository
|
||||
|
||||
@activity.defn(name='validate_train_params')
|
||||
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams:
|
||||
@@ -120,8 +124,8 @@ class Training(SientiaMonitoring):
|
||||
|
||||
This activity orchestrates the ML training pipeline:
|
||||
1. Validate input parameters.
|
||||
2. Train the model via TrainingRepository.
|
||||
3. Perform post-training calculations.
|
||||
2. Prepare data via DataManagerRepository.
|
||||
3. Train the model and compute metrics.
|
||||
|
||||
Args:
|
||||
input_data: Training configuration containing:
|
||||
@@ -130,41 +134,98 @@ class Training(SientiaMonitoring):
|
||||
- train_params (TrainModelParams | dict): Training parameters.
|
||||
|
||||
Returns:
|
||||
dict: Keys `run_name` and `run_dir` when training and saving succeed.
|
||||
dict: Key `run_name` when training and saving succeed.
|
||||
|
||||
Raises:
|
||||
ValueError: If input validation fails.
|
||||
Exception: If training fails (after sending notification).
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
metadata = input_data.get('metadata')
|
||||
train_params = input_data['train_params']
|
||||
|
||||
if isinstance(train_params, dict):
|
||||
train_params = TrainModelParams.from_dict(train_params)
|
||||
# type: ignore[assignment]
|
||||
|
||||
|
||||
model_trained = False
|
||||
model_saved = False
|
||||
metrics_status = 'success'
|
||||
|
||||
try:
|
||||
with self.storage_repository.fetch_file(
|
||||
train_params.bucket_name, train_params.file_name
|
||||
) as uploaded_file:
|
||||
train_result = self.training_repository.train(uploaded_file, train_params)
|
||||
# Download training file bytes from MinIO
|
||||
train_bytes = await self.minio_repository.download_file(
|
||||
object_name=train_params.file_name,
|
||||
bucket=train_params.bucket_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
train_result = self.training_repository.after_train_calculation(
|
||||
train_params, train_result
|
||||
# Download optional validation file bytes from the same bucket
|
||||
val_bytes: bytes | None = None
|
||||
validation_name = getattr(train_params, 'validation_file_name', None)
|
||||
if validation_name is not None:
|
||||
val_bytes = await self.minio_repository.download_file(
|
||||
object_name=validation_name,
|
||||
bucket=train_params.bucket_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
model_trained = True
|
||||
train_result = self.model_repository.save_model(train_result)
|
||||
model_saved = True
|
||||
train_result = self.data_manager_repository.prepare_training_data(
|
||||
train_file_bytes=train_bytes,
|
||||
validation_file_bytes=val_bytes,
|
||||
params=train_params,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
return {
|
||||
'run_name': train_result.run_name,
|
||||
'run_dir': train_result.run_dir,
|
||||
}
|
||||
wrapper = await self.plugin_store.get_model(
|
||||
model_name=train_params.model_name,
|
||||
force_download=False,
|
||||
opt_params={},
|
||||
model_kwargs={},
|
||||
data_model_kwargs={},
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
train_df = pd.concat([train_result.x_train, train_result.y_train], axis=1)
|
||||
val_df = pd.concat([train_result.x_test, train_result.y_test], axis=1)
|
||||
|
||||
wrapper.train(
|
||||
train_data=train_df,
|
||||
val_data=val_df,
|
||||
target=train_params.target_variable,
|
||||
)
|
||||
|
||||
# Generate predictions using the trained wrapper
|
||||
transformed_train, _ = wrapper.transform(train_df)
|
||||
transformed_val, _ = wrapper.transform(val_df)
|
||||
|
||||
y_train_pred_df, _ = wrapper.predict({}, transformed_train)
|
||||
y_val_pred_df, _ = wrapper.predict({}, transformed_val)
|
||||
|
||||
# Use the first column of the prediction DataFrame as the target prediction
|
||||
train_result.y_train_pred = y_train_pred_df.iloc[:, 0]
|
||||
train_result.y_pred = y_val_pred_df.iloc[:, 0]
|
||||
|
||||
train_result = self.data_manager_repository.compute_regression_metrics(
|
||||
train_params,
|
||||
train_result,
|
||||
)
|
||||
|
||||
model_trained = True
|
||||
|
||||
async with self.mlflow_repository.start_run(
|
||||
model_name=train_params.model_name,
|
||||
run_name=None,
|
||||
experiment_name=train_params.experiment_name,
|
||||
tags=None,
|
||||
metadata=metadata,
|
||||
) as run_info:
|
||||
wrapper.store_model(name=train_params.model_name)
|
||||
|
||||
model_saved = True
|
||||
|
||||
return {
|
||||
'run_name': run_info.run_name or run_info.run_id,
|
||||
}
|
||||
except Exception as e: # noqa: BLE001
|
||||
metrics_status = 'error'
|
||||
|
||||
@@ -205,7 +266,6 @@ class Training(SientiaMonitoring):
|
||||
Args:
|
||||
input_data: Cleanup configuration containing:
|
||||
- metadata (dict): Workflow execution metadata.
|
||||
- run_dir (str): Temporary directory to remove.
|
||||
- bucket_name (str): MinIO bucket of the uploaded file.
|
||||
- file_name (str): MinIO object key to delete.
|
||||
|
||||
@@ -213,19 +273,21 @@ class Training(SientiaMonitoring):
|
||||
Exception: If cleanup fails (after sending notification).
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
run_dir = input_data.get('run_dir', '')
|
||||
bucket_name = input_data.get('bucket_name', '')
|
||||
file_name = input_data.get('file_name', '')
|
||||
metrics_status = 'success'
|
||||
|
||||
try:
|
||||
self.model_repository.cleanup_run_directory(run_dir)
|
||||
self.storage_repository.delete_file(bucket_name, file_name)
|
||||
await self.minio_repository.delete_file(
|
||||
object_name=file_name,
|
||||
bucket=bucket_name,
|
||||
metadata=metadata,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001
|
||||
metrics_status = 'error'
|
||||
|
||||
error_msg = (
|
||||
f'Error cleaning up resources - Run directory: {run_dir}, '
|
||||
'Error cleaning up resources - '
|
||||
f'File: {bucket_name}/{file_name}, Error: {str(e)}'
|
||||
)
|
||||
|
||||
|
||||
@@ -121,3 +121,43 @@ def build_minio_config() -> dict[str, Any]:
|
||||
'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')),
|
||||
'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')),
|
||||
}
|
||||
|
||||
|
||||
def build_plugin_store_config() -> dict[str, Any]:
|
||||
"""
|
||||
Build PluginStore configuration from environment variables.
|
||||
|
||||
This function constructs a configuration dictionary for the PluginStore
|
||||
client using environment variables with sensible defaults for local
|
||||
development.
|
||||
|
||||
Environment Variables:
|
||||
STORE_BASE_URL: Base URL of the PluginStore backing Git server
|
||||
(default: http://localhost:3000)
|
||||
STORE_OWNER: Repository owner/organization (default: sientia)
|
||||
STORE_REPO: Repository name (default: model-library-store)
|
||||
STORE_BRANCH: Optional branch name
|
||||
STORE_USERNAME: Optional username for Git HTTP authentication
|
||||
STORE_PASSWORD: Optional password/token for Git HTTP authentication
|
||||
PYPI_SERVER: Optional custom PyPI index URL for runtime installation
|
||||
(default: http://localhost:5000)
|
||||
PYPI_USERNAME: Optional username for PyPI authentication
|
||||
PYPI_PASSWORD: Optional password/token for PyPI authentication
|
||||
|
||||
Returns:
|
||||
dict: PluginStore configuration dictionary with all connector parameters
|
||||
"""
|
||||
|
||||
cache_ttl_seconds = getenv('STORE_CACHE_TTL_SECONDS')
|
||||
return {
|
||||
'base_url': getenv('STORE_BASE_URL', 'http://localhost:3000'),
|
||||
'owner': getenv('STORE_OWNER', 'sientia'),
|
||||
'repo': getenv('STORE_REPO', 'model-library-store'),
|
||||
'username': getenv('STORE_USERNAME'),
|
||||
'password': getenv('STORE_PASSWORD'),
|
||||
'branch': getenv('STORE_BRANCH'),
|
||||
'cache_ttl_seconds': int(cache_ttl_seconds) if cache_ttl_seconds else None,
|
||||
'pypi_index_url': getenv('PYPI_SERVER', 'http://localhost:5000'),
|
||||
'pypi_username': getenv('PYPI_USERNAME'),
|
||||
'pypi_password': getenv('PYPI_PASSWORD'),
|
||||
}
|
||||
|
||||
@@ -33,7 +33,8 @@ class TrainModelParams:
|
||||
use_scaler (bool): Whether to use a scaler for data normalization.
|
||||
include_ar (bool): Whether to include autoregressive variables.
|
||||
bucket_name (str): Name of the MinIO bucket containing training data.
|
||||
file_name (str): Name of the file in the MinIO bucket.
|
||||
file_name (str): Name of the training file in the MinIO bucket.
|
||||
validation_file_name (str | None): Optional name of the validation file in the same MinIO bucket as the training file.
|
||||
line_separator (str): Line separator used in the CSV file.
|
||||
decimal_separator (str): Decimal separator used in the CSV file.
|
||||
date_column (str | None): Name of the date/time column. If set with date_format, the column is parsed as datetime.
|
||||
@@ -66,6 +67,7 @@ class TrainModelParams:
|
||||
include_ar: bool
|
||||
bucket_name: str
|
||||
file_name: str
|
||||
validation_file_name: str | None
|
||||
line_separator: str
|
||||
decimal_separator: str
|
||||
date_column: str | None
|
||||
@@ -122,6 +124,9 @@ class TrainModelParams:
|
||||
include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'),
|
||||
bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
|
||||
file_name=cls._check_none(data.get('file_name'), str, 'file_name'),
|
||||
validation_file_name=cls._check_type(
|
||||
data.get('validation_file_name'), str, 'validation_file_name'
|
||||
),
|
||||
line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
|
||||
decimal_separator=cls._check_none(
|
||||
data.get('decimal_separator'), str, 'decimal_separator'
|
||||
@@ -220,6 +225,7 @@ class TrainModelParams:
|
||||
self._validate_limits()
|
||||
self._validate_required_strings()
|
||||
self._validate_date_format()
|
||||
self._validate_validation_file_name()
|
||||
|
||||
def _validate_numeric_ranges(self) -> None:
|
||||
"""Validate numeric parameters are within acceptable ranges."""
|
||||
@@ -335,3 +341,21 @@ class TrainModelParams:
|
||||
"""Validate date_format is one of the allowed frontend formats when set."""
|
||||
if self.date_format:
|
||||
validate_frontend_date_format(self.date_format)
|
||||
|
||||
def _validate_validation_file_name(self) -> None:
|
||||
"""
|
||||
Validate that validation_file_name, when provided, is not empty or whitespace.
|
||||
|
||||
This field is optional; when present it must point to a valid object key in the
|
||||
same MinIO bucket specified by bucket_name.
|
||||
"""
|
||||
if self.validation_file_name is None:
|
||||
return
|
||||
|
||||
if not isinstance(self.validation_file_name, str):
|
||||
raise TypeError(
|
||||
f'validation_file_name must be a string, got {type(self.validation_file_name).__name__}'
|
||||
)
|
||||
|
||||
if not self.validation_file_name.strip():
|
||||
raise ValueError('validation_file_name cannot be empty or whitespace')
|
||||
|
||||
@@ -2,7 +2,6 @@ from dataclasses import dataclass
|
||||
|
||||
import pandas as pd
|
||||
|
||||
from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
|
||||
|
||||
@@ -12,18 +11,15 @@ class TrainModelResult:
|
||||
A data container for storing the results of a machine learning training process.
|
||||
|
||||
This dataclass encapsulates all outputs from the training pipeline, including
|
||||
the trained model, datasets, evaluation metrics, and paths to generated artifacts.
|
||||
the prepared datasets, evaluation metrics, and paths to generated artifacts.
|
||||
It is used to pass results between activities in the training workflow.
|
||||
|
||||
Attributes:
|
||||
params (TrainModelParams): The parameters used to train the model.
|
||||
process_data (DataPreprocessor): The data preprocessor object used to process the input data.
|
||||
x_train (pd.DataFrame): The training dataset features.
|
||||
x_test (pd.DataFrame): The testing dataset features.
|
||||
y_train (pd.DataFrame): The training dataset target values.
|
||||
y_test (pd.DataFrame): The testing dataset target values.
|
||||
regr (LinearRegressionModel): The trained linear regression model.
|
||||
scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values.
|
||||
y_pred (pd.Series | None): The predicted target values for the testing dataset. Default is None.
|
||||
y_train_pred (pd.Series | None): The predicted target values for the training dataset. Default is None.
|
||||
mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None.
|
||||
@@ -35,17 +31,13 @@ class TrainModelResult:
|
||||
report_path (str | None): The path to the generated HTML report file. Default is None.
|
||||
train_data_path (str | None): The path to the training dataset CSV file. Default is None.
|
||||
test_data_path (str | None): The path to the testing dataset CSV file. Default is None.
|
||||
run_dir (str | None): The path to the run directory containing all artifacts. Default is None.
|
||||
"""
|
||||
|
||||
params: TrainModelParams
|
||||
process_data: DataPreprocessor
|
||||
x_train: pd.DataFrame
|
||||
x_test: pd.DataFrame
|
||||
y_train: pd.Series
|
||||
y_test: pd.Series
|
||||
regr: LinearRegressionModel
|
||||
scaler_dict: dict
|
||||
y_pred: pd.Series | None = None
|
||||
y_train_pred: pd.Series | None = None
|
||||
mse_val: float | None = None
|
||||
@@ -57,4 +49,3 @@ class TrainModelResult:
|
||||
report_path: str | None = None
|
||||
train_data_path: str | None = None
|
||||
test_data_path: str | None = None
|
||||
run_dir: str | None = None
|
||||
|
||||
381
model_manager/utils/repository/data_manager_repository.py
Normal file
381
model_manager/utils/repository/data_manager_repository.py
Normal file
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
Data management repository for the training pipeline.
|
||||
|
||||
This module provides the core data loading and preprocessing logic for the
|
||||
training pipeline, including:
|
||||
- CSV loading from in-memory bytes
|
||||
- datetime parsing and index configuration
|
||||
- optional support filters
|
||||
- train/test split management (when no explicit validation dataset is provided)
|
||||
|
||||
It is intentionally decoupled from any specific model implementation or MLflow
|
||||
integration. Models are trained elsewhere (e.g., via SientiaModel wrappers),
|
||||
and this repository focuses solely on preparing data structures for them.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
|
||||
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
|
||||
|
||||
from model_manager.sientia.metrics import mae, mse, r2
|
||||
from model_manager.sientia.utils import split_train_test
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
|
||||
"""
|
||||
If date_column is set, parse the column as timezone-aware
|
||||
datetime to avoid comparison errors downstream.
|
||||
|
||||
The values are expected to follow the global DATETIME_FORMAT_WITH_TZ
|
||||
pattern defined in sientia_do.temporal.constants.
|
||||
"""
|
||||
if not params.date_column or params.date_column not in data.columns:
|
||||
return data
|
||||
try:
|
||||
data = data.copy()
|
||||
parsed = pd.to_datetime(
|
||||
data[params.date_column],
|
||||
format=DATETIME_FORMAT_WITH_TZ,
|
||||
errors='raise',
|
||||
)
|
||||
data[params.date_column] = parsed
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f'Failed to parse date column "{params.date_column}" with format "{params.date_format}": {e}'
|
||||
) from e
|
||||
return data
|
||||
|
||||
|
||||
def _single_variable_support_mask(
|
||||
data_view: pd.DataFrame,
|
||||
var_col: str,
|
||||
target_variable: str,
|
||||
config: dict,
|
||||
) -> np.ndarray | None:
|
||||
"""Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
|
||||
if var_col not in data_view.columns:
|
||||
return None
|
||||
upper = config.get('upper_line') or config.get('upperLine')
|
||||
lower = config.get('lower_line') or config.get('lowerLine')
|
||||
if not upper or not lower:
|
||||
return None
|
||||
|
||||
x_vals = data_view[var_col].astype(float).to_numpy()
|
||||
y_vals = data_view[target_variable].astype(float).to_numpy()
|
||||
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
|
||||
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
|
||||
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
|
||||
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
|
||||
scale_ratio = y_range / x_range
|
||||
|
||||
b1 = float(upper.get('intercept', 0))
|
||||
deg1 = float(upper.get('angle', 0))
|
||||
b2 = float(lower.get('intercept', 0))
|
||||
deg2 = float(lower.get('angle', 0))
|
||||
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
|
||||
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
|
||||
y1 = m1 * x_vals + b1
|
||||
y2 = m2 * x_vals + b2
|
||||
lower_bound = np.minimum(y1, y2)
|
||||
upper_bound = np.maximum(y1, y2)
|
||||
return (y_vals >= lower_bound) & (y_vals <= upper_bound)
|
||||
|
||||
|
||||
def _apply_support_filters(
|
||||
data_view: pd.DataFrame,
|
||||
target_variable: str,
|
||||
support_filters: dict,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Keep only rows where (var, target) lies between the two guide lines for each variable.
|
||||
|
||||
For each variable in support_filters, the condition is lower(x_var) <= target <= upper(x_var),
|
||||
where lower/upper are the two lines (intercept + slope from angle, scaled by y_range/x_range).
|
||||
Global mask is AND across all variables. Matches DEMO logic in template_01.py.
|
||||
|
||||
Args:
|
||||
data_view: DataFrame after preprocessor transform.
|
||||
target_variable: Name of the target column (y axis).
|
||||
support_filters: Per-variable config with upper_line/lower_line, each {intercept, angle}.
|
||||
|
||||
Returns:
|
||||
data_view filtered to rows satisfying all variable conditions; unchanged if support_filters empty.
|
||||
"""
|
||||
if not support_filters or target_variable not in data_view.columns:
|
||||
return data_view
|
||||
|
||||
combined_keep_mask = np.ones(len(data_view), dtype=bool)
|
||||
n = len(data_view)
|
||||
for var_col, config in support_filters.items():
|
||||
keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
|
||||
if keep_mask is not None and len(keep_mask) == n:
|
||||
combined_keep_mask &= keep_mask
|
||||
|
||||
return data_view.loc[combined_keep_mask]
|
||||
|
||||
|
||||
class DataManagerRepository(SientiaMonitoring):
|
||||
"""
|
||||
Repository for data preparation in the training pipeline.
|
||||
|
||||
This class encapsulates the core logic for preparing ML training data:
|
||||
loading CSV bytes, applying date/index configuration, support filters, and
|
||||
constructing train/test splits (or using an explicit validation dataset).
|
||||
|
||||
Attributes:
|
||||
logger (Logger): Logger instance for observability and debugging
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Logger):
|
||||
"""
|
||||
Initialize DataManagerRepository with logger.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for observability
|
||||
"""
|
||||
SientiaMonitoring.__init__(
|
||||
self,
|
||||
logger=logger,
|
||||
notification_handler=None,
|
||||
metrics_controller=None,
|
||||
)
|
||||
|
||||
def prepare_training_data(
|
||||
self,
|
||||
train_file_bytes: bytes,
|
||||
validation_file_bytes: bytes | None,
|
||||
params: TrainModelParams,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Build TrainModelResult from raw CSV bytes for train (and optional validation) data.
|
||||
|
||||
This method orchestrates the data pipeline:
|
||||
1. Load training data from in-memory bytes
|
||||
2. Optionally load validation data from in-memory bytes
|
||||
3. Parse and configure datetime index
|
||||
4. Apply optional support filters
|
||||
5. Split into train/test sets when no explicit validation dataset is provided
|
||||
|
||||
Args:
|
||||
train_file_bytes: Raw bytes of the training CSV.
|
||||
validation_file_bytes: Raw bytes of the validation CSV, or None when
|
||||
validation should be derived via train/test split.
|
||||
params: Training parameters (TrainModelParams).
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Object containing processed data, train/test splits,
|
||||
and scaler dictionary.
|
||||
|
||||
Raises:
|
||||
ValueError: If transformed data is empty.
|
||||
Exception: If data loading or preprocessing fails.
|
||||
"""
|
||||
try:
|
||||
train_df = pd.read_csv(
|
||||
BytesIO(train_file_bytes),
|
||||
sep=params.line_separator,
|
||||
decimal=params.decimal_separator,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise ValueError(
|
||||
'Failed to load training CSV data from MinIO object. '
|
||||
'Check file encoding, line separator and decimal separator.'
|
||||
) from exc
|
||||
|
||||
train_df = _ensure_date_column_parsed(train_df, params)
|
||||
train_df = self._configure_datetime_index(train_df, params, metadata)
|
||||
|
||||
if params.support_filters:
|
||||
train_df = _apply_support_filters(
|
||||
train_df,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(train_df) <= 0:
|
||||
raise ValueError('Training data view is empty after transformation')
|
||||
|
||||
# Explicit validation dataset path
|
||||
if validation_file_bytes is not None:
|
||||
try:
|
||||
val_df = pd.read_csv(
|
||||
BytesIO(validation_file_bytes),
|
||||
sep=params.line_separator,
|
||||
decimal=params.decimal_separator,
|
||||
)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise ValueError(
|
||||
'Failed to load validation CSV data from MinIO object. '
|
||||
'Check file encoding, line separator and decimal separator.'
|
||||
) from exc
|
||||
|
||||
val_df = _ensure_date_column_parsed(val_df, params)
|
||||
val_df = self._configure_datetime_index(val_df, params, metadata)
|
||||
|
||||
if params.support_filters:
|
||||
val_df = _apply_support_filters(
|
||||
val_df,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(val_df) <= 0:
|
||||
raise ValueError('Validation data view is empty after transformation')
|
||||
|
||||
x_train = pd.DataFrame(train_df[params.variable_columns])
|
||||
y_train = pd.Series(train_df[params.target_variable])
|
||||
x_test = pd.DataFrame(val_df[params.variable_columns])
|
||||
y_test = pd.Series(val_df[params.target_variable])
|
||||
else:
|
||||
# Fallback path: derive validation via train/test split from a single dataset.
|
||||
x_train, x_test, y_train, y_test = split_train_test(
|
||||
pd.DataFrame(train_df[params.variable_columns]),
|
||||
pd.Series(train_df[params.target_variable]),
|
||||
train_size=params.train_size / 100,
|
||||
shuffle=params.shuffle,
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
self.info(
|
||||
f'Data preprocessed and split successfully - experiment run id: {params.experiment_run_id}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
return TrainModelResult(
|
||||
params=params,
|
||||
x_train=x_train,
|
||||
x_test=x_test,
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
)
|
||||
|
||||
def compute_regression_metrics(
|
||||
self,
|
||||
params: TrainModelParams,
|
||||
tmr: TrainModelResult,
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Compute regression metrics for training results.
|
||||
|
||||
This helper mirrors the previous TrainingRepository.after_train_calculation
|
||||
behavior, assuming that y_pred/y_train_pred are already on the correct scale
|
||||
for metric calculation (any scaling is handled inside the model wrapper).
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training.
|
||||
tmr: Training result with y_train, y_test, y_train_pred and y_pred populated.
|
||||
|
||||
Return:
|
||||
Updated TrainModelResult with mse_val, mae_val and r2_val fields populated.
|
||||
"""
|
||||
del params # unused for now, kept for possible future extensions
|
||||
|
||||
tmr.x_train = tmr.x_train.sort_index()
|
||||
tmr.x_test = tmr.x_test.sort_index()
|
||||
tmr.y_train = tmr.y_train.sort_index()
|
||||
tmr.y_test = tmr.y_test.sort_index()
|
||||
|
||||
if tmr.y_pred is not None:
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
if tmr.y_train_pred is not None:
|
||||
tmr.y_train_pred = tmr.y_train_pred.sort_index()
|
||||
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
|
||||
tmr.mse_val = round(
|
||||
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.mae_val = round(
|
||||
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.r2_val = round(
|
||||
r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
return tmr
|
||||
|
||||
def _configure_datetime_index(
|
||||
self,
|
||||
data: pd.DataFrame | None,
|
||||
params: TrainModelParams,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Configure datetime index for the DataFrame.
|
||||
|
||||
Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
|
||||
Prefers params.date_column when set; otherwise looks for common timestamp column names.
|
||||
"""
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Data is None after load_data. '
|
||||
'Check file format, line separator and decimal separator.'
|
||||
)
|
||||
if not isinstance(data, pd.DataFrame):
|
||||
raise TypeError(f'Expected DataFrame, got {type(data).__name__}')
|
||||
|
||||
if isinstance(data.index, pd.DatetimeIndex):
|
||||
self.info('DataFrame already has DatetimeIndex', metadata)
|
||||
return data.sort_index()
|
||||
|
||||
common_timestamp_columns = [
|
||||
'timestamp',
|
||||
'Timestamp',
|
||||
'TIMESTAMP',
|
||||
'date',
|
||||
'Date',
|
||||
'DATE',
|
||||
'DATA',
|
||||
'datetime',
|
||||
'DateTime',
|
||||
]
|
||||
timestamp_columns = ([params.date_column] if params.date_column else []) + [
|
||||
c for c in common_timestamp_columns if c != params.date_column
|
||||
]
|
||||
|
||||
for col in timestamp_columns:
|
||||
if col in data.columns:
|
||||
try:
|
||||
data[col] = pd.to_datetime(data[col])
|
||||
data = data.set_index(col)
|
||||
data = data.sort_index()
|
||||
self.info(f'Configured datetime index from column: {col}', metadata)
|
||||
return data
|
||||
except (ValueError, TypeError) as e:
|
||||
self.warning(f'Failed to convert column {col} to datetime: {e}', metadata)
|
||||
continue
|
||||
|
||||
# If no timestamp column found, check if first column looks like a timestamp
|
||||
first_col = data.columns[0]
|
||||
try:
|
||||
# Try to parse first column as datetime
|
||||
test_values = data[first_col].head(10).dropna()
|
||||
if len(test_values) > 0:
|
||||
pd.to_datetime(test_values)
|
||||
data[first_col] = pd.to_datetime(data[first_col])
|
||||
data = data.set_index(first_col)
|
||||
data = data.sort_index()
|
||||
self.info(f'Configured datetime index from first column: {first_col}', metadata)
|
||||
return data
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
self.warning(
|
||||
'No timestamp column found - some features may not work correctly',
|
||||
metadata,
|
||||
)
|
||||
return data
|
||||
|
||||
@@ -1,471 +0,0 @@
|
||||
"""
|
||||
MLFlow Repository
|
||||
|
||||
This module contains the MLFlowRepository class, which is responsible for
|
||||
handling model training artifacts and MLFlow operations for the Model Manager system.
|
||||
|
||||
It includes methods for generating training reports, managing artifacts,
|
||||
and logging model runs to MLFlow.
|
||||
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
from os import makedirs, path
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
from model_manager.sientia.model_serving import ModelServing # type: ignore[import-untyped]
|
||||
from model_manager.sientia.reports import Reports # type: ignore[import-untyped]
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
# Suppress sklearn FutureWarning about 'squared' deprecation without changing business logic
|
||||
warnings.filterwarnings('ignore', category=FutureWarning, message=".*'squared' is deprecated.*")
|
||||
|
||||
|
||||
class ModelRepository:
|
||||
def __init__(self, url, username, password, logger: Logger):
|
||||
self.model_serving = ModelServing(tracking_uri=url, username=username, password=password)
|
||||
self.logger = logger
|
||||
self.logger.info(f'MLFlow client initialized at {url}')
|
||||
|
||||
def save_model(self, train_result: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Save a trained ML model and its artifacts to MLflow.
|
||||
|
||||
This activity orchestrates the complete model saving pipeline:
|
||||
1. Generates the next run name for the experiment
|
||||
2. Creates and organizes artifacts (reports, data files)
|
||||
3. Logs model, parameters, metrics, and artifacts to MLflow
|
||||
|
||||
Args:
|
||||
input_data: Configuration for model saving operation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- train_result (TrainModelResult): Training result with model and metrics
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated training result with run_name and artifacts
|
||||
|
||||
Raises:
|
||||
Exception: If model saving fails (after sending notification)
|
||||
"""
|
||||
experiment_name = train_result.params.experiment_name
|
||||
train_result.run_name = self._get_next_run_name(experiment_name)
|
||||
train_result = self._generate_artifacts(train_result)
|
||||
self._save_run(train_result)
|
||||
|
||||
self.logger.info(
|
||||
f'Model saved successfully - experiment run id: {train_result.params.experiment_run_id}, '
|
||||
f'experiment name: {experiment_name}, '
|
||||
f'run name: {train_result.run_name}'
|
||||
)
|
||||
|
||||
return train_result
|
||||
|
||||
def cleanup_run_directory(self, run_dir: str) -> None:
|
||||
"""
|
||||
Clean up temporary run directory after model training.
|
||||
|
||||
This activity deletes the temporary directory created during model training
|
||||
and artifact generation. It implements idempotent cleanup to handle cases
|
||||
where the directory may have already been deleted.
|
||||
|
||||
Args:
|
||||
run_dir (str): Path to the run directory to delete
|
||||
"""
|
||||
if not run_dir:
|
||||
self.logger.info('No run directory specified, skipping cleanup')
|
||||
return
|
||||
|
||||
if os.path.exists(run_dir):
|
||||
shutil.rmtree(run_dir)
|
||||
self.logger.info(f'Run directory deleted successfully: {run_dir}')
|
||||
else:
|
||||
self.logger.info(f'Run directory already deleted: {run_dir}')
|
||||
|
||||
def _get_next_run_name(self, experiment_name: str) -> str:
|
||||
"""
|
||||
Generates the next run name for a given experiment.
|
||||
|
||||
Args:
|
||||
experiment_name (str): The name of the experiment for which the next run name is being generated.
|
||||
|
||||
Returns:
|
||||
str: A unique run name in the format "<experiment_name>-<next_run_number>".
|
||||
"""
|
||||
runs = self.model_serving.search_runs_by_name(
|
||||
experiment_names=[experiment_name], order_by=['start_time desc']
|
||||
)
|
||||
|
||||
next_run_number = len(runs) + 1
|
||||
return f'{experiment_name}-{next_run_number}'
|
||||
|
||||
def _generate_artifacts(self, data: TrainModelResult) -> TrainModelResult:
|
||||
"""
|
||||
Generates and organizes artifacts related to the training process, such as reports and data files.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets, model, and parameters.
|
||||
|
||||
Returns:
|
||||
The updated result object with paths to the generated artifacts.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the reports directory or header.html file does not exist.
|
||||
ValueError: If run_name is not set.
|
||||
"""
|
||||
# Validate that run_name is set
|
||||
if not data.run_name:
|
||||
error_msg = 'run_name must be set before generating artifacts'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
reference_data, current_data = self._init_artifacts_data(data)
|
||||
base_path = self._get_reports_directory()
|
||||
|
||||
# Validate that reports directory exists
|
||||
if not path.exists(base_path):
|
||||
error_msg = f'Reports directory does not exist: {base_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg)
|
||||
|
||||
data.run_dir = self._create_run_directory(base_path, data.run_name)
|
||||
header_file_path = path.join(base_path, 'header.html')
|
||||
|
||||
# Validate that header.html exists
|
||||
if not path.exists(header_file_path):
|
||||
error_msg = f'Header file does not exist: {header_file_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg)
|
||||
|
||||
self._setup_run_directory(data.run_dir, header_file_path)
|
||||
return self._generate_report(reference_data, current_data, data)
|
||||
|
||||
def _save_run(self, data: TrainModelResult):
|
||||
"""
|
||||
Logs the details of a machine learning run, including parameters, metrics, models, and artifacts,
|
||||
to the Sientia tracking system.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets, model, parameters,
|
||||
and evaluation metrics.
|
||||
|
||||
Raises:
|
||||
ValueError: If required metrics or artifacts are missing.
|
||||
Exception: If MLflow logging fails for any reason.
|
||||
"""
|
||||
# Validate that required artifacts exist before attempting to log
|
||||
if not data.report_path or not path.exists(data.report_path):
|
||||
error_msg = f'Report file does not exist: {data.report_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if not data.train_data_path or not path.exists(data.train_data_path):
|
||||
error_msg = f'Training data file does not exist: {data.train_data_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if not data.test_data_path or not path.exists(data.test_data_path):
|
||||
error_msg = f'Test data file does not exist: {data.test_data_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Validate that metrics are present
|
||||
if data.mse_val is None or data.r2_val is None or data.mae_val is None:
|
||||
error_msg = 'One or more metrics (MSE, R2, MAE) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Prepare parameters
|
||||
interval_strs = [
|
||||
(str(interval[0]), str(interval[1]))
|
||||
for interval in (data.params.removed_intervals or [])
|
||||
]
|
||||
|
||||
# Set experiment and create run
|
||||
self.model_serving.set_experiment(data.params.experiment_name)
|
||||
|
||||
with self.model_serving.save_experiment(
|
||||
run_name=data.run_name, description=data.params.experiment_name
|
||||
):
|
||||
# Log model parameters
|
||||
self.model_serving.log_param('model_name', data.params.model_name)
|
||||
self.model_serving.log_param(
|
||||
'models_params',
|
||||
{'degree': data.params.degree, 'interaction_only': data.params.interaction_only},
|
||||
)
|
||||
self.model_serving.log_param('target_variable', data.params.target_variable)
|
||||
self.model_serving.log_param('input_variables', data.params.variable_columns)
|
||||
self.model_serving.log_param('nan_treatment', data.params.nan_treatment)
|
||||
self.model_serving.log_param('lag_train', data.params.lag_train)
|
||||
self.model_serving.log_param('lag_transform', data.params.lag_val)
|
||||
static_threshold_value = None
|
||||
if data.params.rem_static_win:
|
||||
static_threshold_value = (
|
||||
data.params.static_threshold if data.params.static_threshold is not None else 1
|
||||
)
|
||||
self.model_serving.log_param('static_threshold', static_threshold_value)
|
||||
self.model_serving.log_param('lower_limits', data.params.low_lim)
|
||||
self.model_serving.log_param('upper_limits', data.params.upp_lim)
|
||||
self.model_serving.log_param('scaler_name', data.params.scaler_name)
|
||||
self.model_serving.log_param('scaler_params', data.scaler_dict)
|
||||
self.model_serving.log_param('include_ar', data.params.include_ar)
|
||||
self.model_serving.log_param('train_size', round(data.params.train_size / 100, 2))
|
||||
self.model_serving.log_param('test_size', round(1 - (data.params.train_size / 100), 2))
|
||||
self.model_serving.log_param('start_date', data.params.start_date)
|
||||
self.model_serving.log_param('end_date', data.params.end_date)
|
||||
self.model_serving.log_param('removed_intervals', interval_strs)
|
||||
self.model_serving.log_param('retrain', False)
|
||||
self.model_serving.log_param('support_filters', data.params.support_filters)
|
||||
|
||||
# Log evaluation metrics
|
||||
self.model_serving.log_metric('MSE', data.mse_val)
|
||||
self.model_serving.log_metric('R2', data.r2_val)
|
||||
self.model_serving.log_metric('MAE', data.mae_val)
|
||||
|
||||
# Log models
|
||||
self.model_serving.log_model(data.process_data, 'data_model')
|
||||
self.model_serving.log_model(data.regr, 'prediction_model')
|
||||
|
||||
# Log artifacts
|
||||
self.model_serving.log_artifact(data.report_path)
|
||||
self.model_serving.log_artifact(data.train_data_path)
|
||||
self.model_serving.log_artifact(data.test_data_path)
|
||||
|
||||
# Log equation artifact if available
|
||||
if data.equation_path and path.exists(data.equation_path):
|
||||
self.model_serving.log_artifact(data.equation_path)
|
||||
|
||||
def _init_artifacts_data(self, data: TrainModelResult) -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""
|
||||
Prepares the reference and current datasets for artifact generation.
|
||||
|
||||
Args:
|
||||
data: The training model result containing the datasets and model.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- reference_data: The training dataset with predictions added.
|
||||
- current_data: The testing dataset with predictions added.
|
||||
|
||||
Raises:
|
||||
ValueError: If training or test datasets are empty or invalid.
|
||||
AttributeError: If required attributes are missing from the data object.
|
||||
"""
|
||||
# Validate that required DataFrames are not empty
|
||||
# Note: x_train, y_train, x_test, y_test, and regr are required fields in TrainModelResult
|
||||
# so we only check if they are empty, not None
|
||||
if data.x_train.empty:
|
||||
error_msg = 'Training features (x_train) are empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_train.empty:
|
||||
error_msg = 'Training target (y_train) is empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.x_test.empty:
|
||||
error_msg = 'Test features (x_test) are empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_test.empty:
|
||||
error_msg = 'Test target (y_test) is empty'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Validate that predictions exist (y_pred is optional, so check for None)
|
||||
if data.y_pred is None:
|
||||
error_msg = 'Test predictions (y_pred) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if data.y_train_pred is None:
|
||||
error_msg = 'Training predictions (y_train_pred) are None'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Prepare reference data (training set)
|
||||
reference_data = pd.concat([data.x_train, data.y_train], axis=1)
|
||||
reference_data = reference_data.rename(columns={data.params.target_variable: 'target'})
|
||||
# Use pre-calculated predictions (calculated before denormalization to avoid overflow)
|
||||
reference_data['prediction'] = data.y_train_pred
|
||||
|
||||
# Prepare current data (test set)
|
||||
current_data = pd.concat([data.x_test, data.y_test], axis=1)
|
||||
current_data = current_data.rename(columns={data.params.target_variable: 'target'})
|
||||
current_data['prediction'] = data.y_pred
|
||||
|
||||
return reference_data, current_data
|
||||
|
||||
def _create_run_directory(self, base_path: str, run_name: str) -> str:
|
||||
"""
|
||||
Creates a directory inside the 'reports' folder with the run name and a timestamp.
|
||||
|
||||
Uses microsecond precision in timestamp to minimize collision probability
|
||||
in high-concurrency scenarios.
|
||||
|
||||
Args:
|
||||
base_path (str): The path to the 'reports' folder.
|
||||
run_name (str): The name of the run.
|
||||
|
||||
Returns:
|
||||
str: The path to the created directory.
|
||||
|
||||
Raises:
|
||||
PermissionError: If there are insufficient permissions to create the directory.
|
||||
OSError: If directory creation fails for any other reason.
|
||||
"""
|
||||
# Use microsecond precision to reduce collision probability
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S_%f')
|
||||
run_dir = path.join(base_path, 'temp', f'{run_name}_{timestamp}')
|
||||
|
||||
try:
|
||||
makedirs(run_dir, exist_ok=True)
|
||||
return run_dir
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when creating directory: {run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to create directory {run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _setup_run_directory(self, run_dir: str, header_file_path: str):
|
||||
"""
|
||||
Creates empty files and copies a header file into the specified run directory.
|
||||
|
||||
Note: Lock removed as each run has its own unique directory, so no synchronization
|
||||
is needed between different runs. File operations within the same directory are
|
||||
atomic at the OS level.
|
||||
|
||||
Args:
|
||||
run_dir (str): The path to the run directory where the files will be created.
|
||||
header_file_path (str): The path to the header.html file to be copied.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the header file does not exist.
|
||||
PermissionError: If there are insufficient permissions to create files.
|
||||
OSError: If file creation or copying fails for any other reason.
|
||||
"""
|
||||
empty_files = ['data_drift.html', 'data_quality.html', 'regression.html']
|
||||
|
||||
try:
|
||||
# Create empty placeholder files
|
||||
for file_name in empty_files:
|
||||
file_path = path.join(run_dir, file_name)
|
||||
with open(file_path, 'w'):
|
||||
pass # Create empty file
|
||||
|
||||
# Copy header file to run directory
|
||||
header_dest = path.join(run_dir, 'header.html')
|
||||
shutil.copy(header_file_path, header_dest)
|
||||
except FileNotFoundError as e:
|
||||
error_msg = f'Header file not found: {header_file_path}'
|
||||
self.logger.error(error_msg)
|
||||
raise FileNotFoundError(error_msg) from e
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when setting up directory: {run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to setup run directory {run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _generate_report(
|
||||
self, reference_data: pd.DataFrame, current_data: pd.DataFrame, data: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Generates a comprehensive report summarizing data quality, data drift, and regression analysis.
|
||||
|
||||
Args:
|
||||
reference_data (pd.DataFrame): The training dataset with predictions added.
|
||||
current_data (pd.DataFrame): The testing dataset with predictions added.
|
||||
data: The training model result containing the datasets, model, and parameters.
|
||||
|
||||
Returns:
|
||||
The updated result object with paths to the generated report and data files.
|
||||
|
||||
Raises:
|
||||
ValueError: If data conversion to float64 fails or DataFrames are invalid.
|
||||
PermissionError: If there are insufficient permissions to write files.
|
||||
OSError: If file writing fails for any other reason.
|
||||
"""
|
||||
try:
|
||||
# Convert data to float64 for report generation
|
||||
# This may raise ValueError if data contains non-numeric values
|
||||
reference_data_float = reference_data.astype(np.float64)
|
||||
current_data_float = current_data.astype(np.float64)
|
||||
|
||||
# Initialize report generator
|
||||
report = Reports(
|
||||
reference_data=reference_data_float,
|
||||
current_data=current_data_float,
|
||||
base_path=data.run_dir,
|
||||
)
|
||||
|
||||
# Generate report sections
|
||||
report.add_data_quality_section(columns=data.params.variable_columns + ['target'])
|
||||
report.add_data_drift_section(columns=data.params.variable_columns + ['target'])
|
||||
report.add_regression_section()
|
||||
|
||||
# Validate that run_dir is set (should be set by _create_run_directory)
|
||||
if not data.run_dir:
|
||||
error_msg = 'run_dir is not set after directory creation'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
# Save HTML report
|
||||
data.report_path = path.join(data.run_dir, 'report.html')
|
||||
report.save_all_sections_html(data.report_path)
|
||||
|
||||
# Save training data CSV
|
||||
data.train_data_path = path.join(data.run_dir, 'train_data.csv')
|
||||
reference_data.to_csv(data.train_data_path, index=False)
|
||||
|
||||
# Save test data CSV
|
||||
data.test_data_path = path.join(data.run_dir, 'test_data.csv')
|
||||
current_data.to_csv(data.test_data_path, index=False)
|
||||
|
||||
# Save equation as JSON
|
||||
if data.equation is not None:
|
||||
data.equation_path = path.join(data.run_dir, 'model_equation.json')
|
||||
with open(data.equation_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(data.equation, f, indent=2, ensure_ascii=False)
|
||||
|
||||
return data
|
||||
except ValueError as e:
|
||||
error_msg = f'Failed to convert data to float64 for report generation: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise ValueError(error_msg) from e
|
||||
except PermissionError as e:
|
||||
error_msg = f'Permission denied when writing report files to: {data.run_dir}'
|
||||
self.logger.error(error_msg)
|
||||
raise PermissionError(error_msg) from e
|
||||
except OSError as e:
|
||||
error_msg = f'Failed to generate report in {data.run_dir}: {str(e)}'
|
||||
self.logger.error(error_msg)
|
||||
raise OSError(error_msg) from e
|
||||
|
||||
def _get_reports_directory(self) -> str:
|
||||
"""
|
||||
Get the absolute path to the reports directory.
|
||||
|
||||
Returns:
|
||||
str: Absolute path to model_manager/reports directory.
|
||||
"""
|
||||
# Get the directory where this file is located (model_manager/utils/repository/)
|
||||
current_file_dir = path.dirname(path.abspath(__file__))
|
||||
# Navigate up to model_manager/ and then to reports/
|
||||
model_manager_dir = path.dirname(path.dirname(current_file_dir))
|
||||
reports_dir = path.join(model_manager_dir, 'reports')
|
||||
return reports_dir
|
||||
@@ -1,164 +0,0 @@
|
||||
from io import BytesIO
|
||||
|
||||
import boto3 # type: ignore[import-untyped]
|
||||
from botocore.config import Config # type: ignore[import-untyped]
|
||||
from sientia_do.observability.logger import Logger
|
||||
|
||||
|
||||
class StorageRepository:
|
||||
"""
|
||||
MinIO (S3-compatible) storage activities for file operations.
|
||||
|
||||
This class provides activities for interacting with MinIO object storage,
|
||||
including file download and deletion operations. It handles authentication,
|
||||
connection management, and comprehensive error handling.
|
||||
|
||||
The class implements best practices for S3/MinIO operations:
|
||||
- Connection reuse (boto3 client is thread-safe)
|
||||
- Automatic retry with exponential backoff
|
||||
- Comprehensive error handling and logging
|
||||
- Notification integration for critical errors
|
||||
|
||||
Attributes:
|
||||
endpoint_url (str): MinIO server endpoint URL
|
||||
access_key (str): MinIO access key ID
|
||||
secret_key (str): MinIO secret access key
|
||||
region (str): MinIO region name
|
||||
use_ssl (bool): Whether to use SSL/TLS for connections
|
||||
minio_client: Boto3 S3 client configured for MinIO
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
endpoint_url: str,
|
||||
access_key: str,
|
||||
secret_key: str,
|
||||
region: str,
|
||||
use_ssl: bool,
|
||||
max_retry_attempts: int,
|
||||
retry_mode: str,
|
||||
connect_timeout: int,
|
||||
read_timeout: int,
|
||||
logger: Logger,
|
||||
):
|
||||
"""
|
||||
Initialize a reusable MinIO client with retry configuration.
|
||||
|
||||
Args:
|
||||
endpoint_url: MinIO server endpoint URL (e.g., http://localhost:9000).
|
||||
access_key: MinIO access key ID for authentication.
|
||||
secret_key: MinIO secret access key for authentication.
|
||||
region: MinIO region name (e.g., us-east-1).
|
||||
use_ssl: Whether to use SSL/TLS for connections.
|
||||
max_retry_attempts: Maximum number of retry attempts (e.g., 3).
|
||||
retry_mode: Retry policy to apply (standard, legacy, adaptive).
|
||||
connect_timeout: Connection timeout in seconds.
|
||||
read_timeout: Read timeout in seconds.
|
||||
logger: Logger used for observability.
|
||||
"""
|
||||
self.endpoint_url = endpoint_url
|
||||
self.access_key = access_key
|
||||
self.secret_key = secret_key
|
||||
self.region = region
|
||||
self.use_ssl = use_ssl
|
||||
self.max_retry_attempts = max_retry_attempts
|
||||
self.retry_mode = retry_mode
|
||||
self.connect_timeout = connect_timeout
|
||||
self.read_timeout = read_timeout
|
||||
self.logger = logger
|
||||
|
||||
boto_config = Config(
|
||||
region_name=region,
|
||||
retries={
|
||||
'max_attempts': max_retry_attempts,
|
||||
'mode': retry_mode,
|
||||
},
|
||||
connect_timeout=connect_timeout,
|
||||
read_timeout=read_timeout,
|
||||
)
|
||||
|
||||
self.minio_client = boto3.client(
|
||||
's3',
|
||||
endpoint_url=endpoint_url,
|
||||
aws_access_key_id=access_key,
|
||||
aws_secret_access_key=secret_key,
|
||||
config=boto_config,
|
||||
use_ssl=use_ssl,
|
||||
)
|
||||
|
||||
self.logger.info(f'MinIO client initialized at {endpoint_url}')
|
||||
|
||||
def close(self) -> None:
|
||||
self.minio_client.close()
|
||||
self.logger.info('MinIO client closed')
|
||||
|
||||
def fetch_file(self, bucket_name: str, file_name: str) -> BytesIO:
|
||||
"""
|
||||
Fetch an object from MinIO and return its contents as `BytesIO`.
|
||||
|
||||
Args:
|
||||
bucket_name: MinIO bucket where the object resides.
|
||||
file_name: Object key to download inside the bucket.
|
||||
|
||||
Returns:
|
||||
BytesIO: File-like stream containing the downloaded bytes.
|
||||
|
||||
Raises:
|
||||
OSError: If the download fails (network, permissions, missing key, etc.).
|
||||
"""
|
||||
response = self.minio_client.get_object(Bucket=bucket_name, Key=file_name)
|
||||
|
||||
with response['Body'] as body:
|
||||
file_content = body.read()
|
||||
|
||||
file_size = len(file_content)
|
||||
|
||||
self.logger.info(
|
||||
f'File fetched successfully: {bucket_name}/{file_name} ({file_size} bytes)'
|
||||
)
|
||||
|
||||
return BytesIO(file_content)
|
||||
|
||||
def delete_file(self, bucket_name: str, file_name: str) -> None:
|
||||
"""
|
||||
Remove an object from MinIO storage.
|
||||
|
||||
Args:
|
||||
bucket_name: Bucket that contains the object.
|
||||
file_name: Object key to delete.
|
||||
"""
|
||||
self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
|
||||
self.logger.info(f'File deleted successfully: {bucket_name}/{file_name}')
|
||||
|
||||
def list_bucket_objects(self, bucket_name: str, max_keys: int = 1000) -> list[str]:
|
||||
"""
|
||||
List objects in a MinIO bucket.
|
||||
|
||||
This method uses the MinIO/S3 list_objects_v2 API to retrieve objects
|
||||
from the specified bucket. This is optimized for cleanup operations
|
||||
by using configurable pagination.
|
||||
|
||||
Args:
|
||||
bucket_name: Name of the bucket to list objects from.
|
||||
max_keys: Maximum number of keys per page (default: 1000).
|
||||
|
||||
Returns:
|
||||
List[str]: List of object keys (file names).
|
||||
"""
|
||||
# Use list_objects_v2 for efficient pagination
|
||||
paginator = self.minio_client.get_paginator('list_objects_v2')
|
||||
|
||||
pages = paginator.paginate(Bucket=bucket_name, MaxKeys=max_keys)
|
||||
|
||||
objects = []
|
||||
total_count = 0
|
||||
|
||||
for page in pages:
|
||||
if 'Contents' in page:
|
||||
for obj in page['Contents']:
|
||||
objects.append(obj['Key'])
|
||||
total_count += 1
|
||||
|
||||
self.logger.info(f'Listed {total_count} objects from bucket {bucket_name}')
|
||||
|
||||
return objects
|
||||
@@ -1,519 +0,0 @@
|
||||
"""
|
||||
Training repository for ML model training operations.
|
||||
|
||||
This module provides the core training logic for machine learning models,
|
||||
including data preprocessing, model training, and post-training calculations.
|
||||
Migrated from laborious/utils/train_model_utils.py.
|
||||
"""
|
||||
|
||||
from io import BytesIO
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from sientia_do.observability.logger import Logger
|
||||
from sientia_do.operations.df_preprocessor import load_data
|
||||
from sientia_do.operations.normalization import MinMaxScaler, Z_Scaler
|
||||
|
||||
from model_manager.sientia.metrics import mae, mse, r2
|
||||
from model_manager.sientia.models import (
|
||||
DataPreprocessor,
|
||||
LinearRegressionModel,
|
||||
)
|
||||
from model_manager.sientia.models import (
|
||||
_frontend_date_format_to_strftime as _frontend_format_to_strftime,
|
||||
)
|
||||
from model_manager.sientia.utils import split_train_test
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.models.train_model_result import TrainModelResult
|
||||
|
||||
|
||||
def _ensure_date_column_parsed(data: pd.DataFrame, params: TrainModelParams) -> pd.DataFrame:
|
||||
"""If date_column and date_format are set, parse the column as datetime to avoid comparison errors downstream."""
|
||||
if not params.date_column or not params.date_format or params.date_column not in data.columns:
|
||||
return data
|
||||
try:
|
||||
python_fmt = _frontend_format_to_strftime(params.date_format)
|
||||
data = data.copy()
|
||||
data[params.date_column] = pd.to_datetime(
|
||||
data[params.date_column], format=python_fmt, errors='coerce'
|
||||
)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f'Failed to parse date column "{params.date_column}" with format "{params.date_format}": {e}'
|
||||
) from e
|
||||
return data
|
||||
|
||||
|
||||
def _single_variable_support_mask(
|
||||
data_view: pd.DataFrame,
|
||||
var_col: str,
|
||||
target_variable: str,
|
||||
config: dict,
|
||||
) -> np.ndarray | None:
|
||||
"""Compute keep mask for one variable's support lines; None if config is invalid or skipped."""
|
||||
if var_col not in data_view.columns:
|
||||
return None
|
||||
upper = config.get('upper_line') or config.get('upperLine')
|
||||
lower = config.get('lower_line') or config.get('lowerLine')
|
||||
if not upper or not lower:
|
||||
return None
|
||||
|
||||
x_vals = data_view[var_col].astype(float).to_numpy()
|
||||
y_vals = data_view[target_variable].astype(float).to_numpy()
|
||||
xmin, xmax = float(np.nanmin(x_vals)), float(np.nanmax(x_vals))
|
||||
ymin, ymax = float(np.nanmin(y_vals)), float(np.nanmax(y_vals))
|
||||
x_range = (xmax - xmin) if (xmax - xmin) != 0 else 1.0
|
||||
y_range = (ymax - ymin) if (ymax - ymin) != 0 else 1.0
|
||||
scale_ratio = y_range / x_range
|
||||
|
||||
b1 = float(upper.get('intercept', 0))
|
||||
deg1 = float(upper.get('angle', 0))
|
||||
b2 = float(lower.get('intercept', 0))
|
||||
deg2 = float(lower.get('angle', 0))
|
||||
m1 = np.tan(np.deg2rad(deg1)) * scale_ratio
|
||||
m2 = np.tan(np.deg2rad(deg2)) * scale_ratio
|
||||
y1 = m1 * x_vals + b1
|
||||
y2 = m2 * x_vals + b2
|
||||
lower_bound = np.minimum(y1, y2)
|
||||
upper_bound = np.maximum(y1, y2)
|
||||
return (y_vals >= lower_bound) & (y_vals <= upper_bound)
|
||||
|
||||
|
||||
def _apply_support_filters(
|
||||
data_view: pd.DataFrame,
|
||||
target_variable: str,
|
||||
support_filters: dict,
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Keep only rows where (var, target) lies between the two guide lines for each variable.
|
||||
|
||||
For each variable in support_filters, the condition is lower(x_var) <= target <= upper(x_var),
|
||||
where lower/upper are the two lines (intercept + slope from angle, scaled by y_range/x_range).
|
||||
Global mask is AND across all variables. Matches DEMO logic in template_01.py.
|
||||
|
||||
Args:
|
||||
data_view: DataFrame after preprocessor transform.
|
||||
target_variable: Name of the target column (y axis).
|
||||
support_filters: Per-variable config with upper_line/lower_line, each {intercept, angle}.
|
||||
|
||||
Returns:
|
||||
data_view filtered to rows satisfying all variable conditions; unchanged if support_filters empty.
|
||||
"""
|
||||
if not support_filters or target_variable not in data_view.columns:
|
||||
return data_view
|
||||
|
||||
combined_keep_mask = np.ones(len(data_view), dtype=bool)
|
||||
n = len(data_view)
|
||||
for var_col, config in support_filters.items():
|
||||
keep_mask = _single_variable_support_mask(data_view, var_col, target_variable, config)
|
||||
if keep_mask is not None and len(keep_mask) == n:
|
||||
combined_keep_mask &= keep_mask
|
||||
|
||||
return data_view.loc[combined_keep_mask]
|
||||
|
||||
|
||||
class TrainingRepository:
|
||||
"""
|
||||
Repository for machine learning model training operations.
|
||||
|
||||
This class encapsulates the core logic for training ML models, migrated from
|
||||
laborious/utils/train_model_utils.py. Follows the same pattern as MLFlowRepository
|
||||
with instance methods and logger integration.
|
||||
|
||||
Attributes:
|
||||
logger (Logger): Logger instance for observability and debugging
|
||||
"""
|
||||
|
||||
def __init__(self, logger: Logger):
|
||||
"""
|
||||
Initialize TrainingRepository with logger.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for observability
|
||||
"""
|
||||
self.logger = logger
|
||||
|
||||
def train(self, uploaded_file: BytesIO, params: TrainModelParams) -> TrainModelResult:
|
||||
"""
|
||||
Train a machine learning model using the provided file and parameters.
|
||||
|
||||
This method orchestrates the training pipeline:
|
||||
1. Load data from BytesIO file
|
||||
2. Initialize and fit data preprocessor
|
||||
3. Transform data and validate
|
||||
4. Split into train/test sets
|
||||
5. Initialize scaler dictionary
|
||||
6. Train LinearRegression model
|
||||
|
||||
Args:
|
||||
uploaded_file: BytesIO object containing training data (CSV format)
|
||||
params: Training parameters (TrainModelParams)
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Object containing trained model, processed data,
|
||||
train/test splits, and scaler dictionary
|
||||
|
||||
Raises:
|
||||
ValueError: If transformed data is empty
|
||||
Exception: If data loading, preprocessing, or training fails
|
||||
"""
|
||||
data = load_data(uploaded_file, params.line_separator, params.decimal_separator)
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Failed to load CSV data: load_data returned None. '
|
||||
'Check file encoding, line separator and 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)
|
||||
|
||||
if params.support_filters:
|
||||
data_view = _apply_support_filters(
|
||||
data_view,
|
||||
params.target_variable,
|
||||
params.support_filters,
|
||||
)
|
||||
|
||||
if len(data_view) <= 0:
|
||||
raise ValueError('Data view is empty after transformation')
|
||||
|
||||
x_train, x_test, y_train, y_test = split_train_test(
|
||||
data_view[params.variable_columns],
|
||||
data_view[params.target_variable],
|
||||
train_size=params.train_size / 100,
|
||||
shuffle=params.shuffle,
|
||||
random_state=42,
|
||||
)
|
||||
|
||||
data_train = pd.concat([x_train, y_train], axis=1)
|
||||
scaler_dict = self._init_scaler_dict(process_data, params)
|
||||
|
||||
regr = LinearRegressionModel(
|
||||
target_variable=params.target_variable,
|
||||
variable_columns=params.variable_columns,
|
||||
degree=params.degree,
|
||||
interaction_only=params.interaction_only,
|
||||
)
|
||||
|
||||
regr.fit(data_train)
|
||||
self.logger.info(
|
||||
f'Model trained successfully - experiment run id: {params.experiment_run_id}'
|
||||
)
|
||||
|
||||
return TrainModelResult(
|
||||
params=params,
|
||||
process_data=process_data,
|
||||
x_train=x_train,
|
||||
x_test=x_test,
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=regr,
|
||||
scaler_dict=scaler_dict,
|
||||
)
|
||||
|
||||
def after_train_calculation(
|
||||
self, params: TrainModelParams, tmr: TrainModelResult
|
||||
) -> TrainModelResult:
|
||||
"""
|
||||
Perform post-training calculations: predictions, denormalization, and metrics.
|
||||
|
||||
This method completes the training pipeline by:
|
||||
1. Making predictions on test set
|
||||
2. Denormalizing all data (if scaler was used)
|
||||
3. Reordering data by index
|
||||
4. Calculating evaluation metrics (MSE, MAE, R²)
|
||||
|
||||
Args:
|
||||
params: Training parameters used during model training
|
||||
tmr: Result object from training
|
||||
|
||||
Returns:
|
||||
TrainModelResult: Updated result with predictions, denormalized data,
|
||||
and metrics (mse_val, mae_val, r2_val)
|
||||
"""
|
||||
# Calculate predictions BEFORE denormalization (important for polynomial models)
|
||||
y_pred_array = tmr.regr.predict(tmr.x_test)
|
||||
y_train_pred_array = tmr.regr.predict(tmr.x_train)
|
||||
|
||||
if params.use_scaler:
|
||||
scaler = tmr.process_data.get_scaler()
|
||||
|
||||
# If using custom scaler with denormalize_* helpers
|
||||
if hasattr(scaler, 'denormalize_single_input'):
|
||||
for col in params.variable_columns:
|
||||
tmr.x_train[col] = scaler.denormalize_single_input(tmr.x_train[col], col)
|
||||
tmr.x_test[col] = scaler.denormalize_single_input(tmr.x_test[col], col)
|
||||
|
||||
tmr.y_train = scaler.denormalize_single_input(tmr.y_train, params.target_variable)
|
||||
tmr.y_test = scaler.denormalize_single_input(tmr.y_test, params.target_variable)
|
||||
y_pred_array = scaler.denormalize_predictions(y_pred_array, params.target_variable)
|
||||
y_train_pred_array = scaler.denormalize_predictions(
|
||||
y_train_pred_array, params.target_variable
|
||||
)
|
||||
else:
|
||||
# Fallback for sklearn StandardScaler: only inverse-transform features
|
||||
feature_cols = getattr(
|
||||
tmr.process_data, 'feature_names_order', params.variable_columns
|
||||
)
|
||||
# Ensure columns are in the same order used during fit
|
||||
x_train_features = tmr.x_train[feature_cols]
|
||||
x_test_features = tmr.x_test[feature_cols]
|
||||
|
||||
tmr.x_train[feature_cols] = scaler.inverse_transform(x_train_features)
|
||||
tmr.x_test[feature_cols] = scaler.inverse_transform(x_test_features)
|
||||
# Target was not scaled with StandardScaler in preprocessing; leave y as-is
|
||||
|
||||
tmr.y_pred = pd.Series(y_pred_array, index=tmr.y_test.index)
|
||||
tmr.y_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
tmr.y_train_pred = pd.Series(y_train_pred_array, index=tmr.y_train.index)
|
||||
tmr.y_train_pred.name = f'{params.target_variable}_pred'
|
||||
|
||||
tmr.x_train = tmr.x_train.sort_index()
|
||||
tmr.x_test = tmr.x_test.sort_index()
|
||||
tmr.y_train = tmr.y_train.sort_index()
|
||||
tmr.y_test = tmr.y_test.sort_index()
|
||||
tmr.y_pred = tmr.y_pred.sort_index()
|
||||
tmr.y_train_pred = tmr.y_train_pred.sort_index()
|
||||
|
||||
assert tmr.y_pred is not None, 'y_pred should be set at this point'
|
||||
|
||||
tmr.mse_val = round(
|
||||
mse(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.mae_val = round(
|
||||
mae(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)),
|
||||
2,
|
||||
)
|
||||
|
||||
tmr.r2_val = round(r2(tmr.y_test.astype(np.float64), tmr.y_pred.astype(np.float64)), 2)
|
||||
|
||||
# Extract model equation
|
||||
tmr.equation = self._extract_model_equation(tmr.regr, params)
|
||||
|
||||
self.logger.info(
|
||||
f'Model metrics calculated successfully - experiment run id: {params.experiment_run_id}'
|
||||
)
|
||||
return tmr
|
||||
|
||||
def _init_scaler_dict(self, process_data: DataPreprocessor, params: TrainModelParams) -> dict:
|
||||
"""
|
||||
Initialize dictionary containing scaling parameters for features and target.
|
||||
|
||||
This method extracts scaling parameters from the fitted scaler to enable
|
||||
denormalization of predictions and debugging of the normalization process.
|
||||
|
||||
Args:
|
||||
process_data: Fitted DataPreprocessor object with scaler
|
||||
params: Training parameters including scaler configuration
|
||||
|
||||
Returns:
|
||||
dict: Scaling parameters for each feature and target variable.
|
||||
Structure depends on scaler type:
|
||||
- MinMaxScaler: {'feature': {'min': float, 'max': float}, ...}
|
||||
- Z_Scaler: Dictionary from scaler.create_dict()
|
||||
- Empty dict: If no scaler is used
|
||||
|
||||
Raises:
|
||||
AttributeError: If scaler doesn't have expected attributes
|
||||
"""
|
||||
scaler_dict = {}
|
||||
|
||||
if params.use_scaler:
|
||||
scaler = process_data.get_scaler()
|
||||
|
||||
if isinstance(scaler, MinMaxScaler):
|
||||
# Extract min/max for each feature
|
||||
for i, col in enumerate(params.variable_columns):
|
||||
scaler_dict[col] = {'min': scaler.x_min[i], 'max': scaler.x_max[i]}
|
||||
|
||||
# Extract min/max for target variable
|
||||
scaler_dict[params.target_variable] = {
|
||||
'min': scaler.y_min,
|
||||
'max': scaler.y_max,
|
||||
}
|
||||
|
||||
elif isinstance(scaler, Z_Scaler):
|
||||
scaler_dict = scaler.create_dict()
|
||||
|
||||
return scaler_dict
|
||||
|
||||
def _get_static_threshold(self, params: TrainModelParams) -> int | None:
|
||||
"""
|
||||
Get the static threshold value based on parameters.
|
||||
|
||||
Args:
|
||||
params: Training parameters containing static window configuration
|
||||
|
||||
Returns:
|
||||
int | None: Static threshold value (1-1000) if rem_static_win is True, None otherwise
|
||||
"""
|
||||
if not params.rem_static_win:
|
||||
return None
|
||||
return params.static_threshold if params.static_threshold is not None else 1
|
||||
|
||||
def _init_data_preprocessor(self, params: TrainModelParams) -> DataPreprocessor:
|
||||
"""
|
||||
Initialize DataPreprocessor with training parameters.
|
||||
|
||||
Args:
|
||||
params: Training parameters containing preprocessor configuration
|
||||
|
||||
Returns:
|
||||
DataPreprocessor: Configured preprocessor ready for fitting
|
||||
"""
|
||||
# Convert removed_intervals to list of tuples if needed
|
||||
removed_intervals = None
|
||||
if params.removed_intervals:
|
||||
removed_intervals = [
|
||||
(interval[0], interval[1]) if isinstance(interval, (list, tuple)) else interval
|
||||
for interval in params.removed_intervals
|
||||
]
|
||||
|
||||
return DataPreprocessor(
|
||||
target_variable=params.target_variable,
|
||||
input_columns=params.variable_columns,
|
||||
nan_treatment=params.nan_treatment,
|
||||
lag_train=params.lag_train,
|
||||
lag_transform=params.lag_val,
|
||||
start_date=params.start_date,
|
||||
end_date=params.end_date,
|
||||
date_format=params.date_format,
|
||||
removed_intervals=removed_intervals,
|
||||
static_threshold=self._get_static_threshold(params),
|
||||
low_lim=params.low_lim,
|
||||
upp_lim=params.upp_lim,
|
||||
scaler_name=params.scaler_name,
|
||||
scaler_params={} if params.use_scaler else None,
|
||||
ar_var=params.target_variable if params.include_ar else None,
|
||||
)
|
||||
|
||||
def _extract_model_equation(
|
||||
self, regr: LinearRegressionModel, params: TrainModelParams
|
||||
) -> dict:
|
||||
"""
|
||||
Extract the linear regression equation coefficients and create equation metadata.
|
||||
|
||||
This method extracts the coefficients and intercept from the trained model
|
||||
and creates a structured dictionary containing the equation information
|
||||
for serialization as JSON artifact.
|
||||
|
||||
Args:
|
||||
regr: Trained LinearRegressionModel object
|
||||
params: Training parameters containing variable information
|
||||
|
||||
Returns:
|
||||
dict: Equation metadata containing:
|
||||
- target_variable: Name of the target variable
|
||||
- coefficients: Dictionary mapping variable names to coefficients
|
||||
- intercept: Model intercept value
|
||||
- equation_string: Human-readable equation string
|
||||
- latex_equation: LaTeX formatted equation
|
||||
"""
|
||||
coefficients = regr.regr.coef_
|
||||
intercept = regr.regr.intercept_
|
||||
|
||||
# Get feature names - for polynomial models, use poly_feature_names
|
||||
if params.degree > 1 and regr.poly_feature_names:
|
||||
feature_names = regr.poly_feature_names
|
||||
else:
|
||||
feature_names = params.variable_columns
|
||||
|
||||
# Create coefficients dictionary
|
||||
coefficients_dict = {}
|
||||
for i, var in enumerate(feature_names):
|
||||
if i < len(coefficients):
|
||||
coefficients_dict[var] = float(coefficients[i])
|
||||
|
||||
# Create equation string
|
||||
equation_parts = [f'{coef:.6f} * {var}' for var, coef in coefficients_dict.items()]
|
||||
equation_string = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(
|
||||
equation_parts
|
||||
)
|
||||
|
||||
# Create LaTeX equation
|
||||
latex_parts = [f'{coef:.6f} \\cdot {var}' for var, coef in coefficients_dict.items()]
|
||||
latex_equation = f'{params.target_variable} = {intercept:.6f} + ' + ' + '.join(latex_parts)
|
||||
|
||||
return {
|
||||
'target_variable': params.target_variable,
|
||||
'coefficients': coefficients_dict,
|
||||
'intercept': float(intercept),
|
||||
'equation_string': equation_string,
|
||||
'latex_equation': latex_equation,
|
||||
'model_type': params.model_name,
|
||||
'degree': params.degree,
|
||||
'interaction_only': params.interaction_only,
|
||||
'original_features': params.variable_columns,
|
||||
}
|
||||
|
||||
def _configure_datetime_index(
|
||||
self, data: pd.DataFrame | None, params: TrainModelParams
|
||||
) -> pd.DataFrame:
|
||||
"""
|
||||
Configure datetime index for the DataFrame.
|
||||
|
||||
Guards against None to avoid 'NoneType' object has no attribute 'index' downstream.
|
||||
Prefers params.date_column when set; otherwise looks for common timestamp column names.
|
||||
"""
|
||||
if data is None:
|
||||
raise ValueError(
|
||||
'Data is None after load_data. '
|
||||
'Check file format, line separator and decimal separator.'
|
||||
)
|
||||
if not isinstance(data, pd.DataFrame):
|
||||
raise TypeError(f'Expected DataFrame, got {type(data).__name__}')
|
||||
|
||||
if isinstance(data.index, pd.DatetimeIndex):
|
||||
self.logger.info('DataFrame already has DatetimeIndex')
|
||||
return data.sort_index()
|
||||
|
||||
common_timestamp_columns = [
|
||||
'timestamp',
|
||||
'Timestamp',
|
||||
'TIMESTAMP',
|
||||
'date',
|
||||
'Date',
|
||||
'DATE',
|
||||
'DATA',
|
||||
'datetime',
|
||||
'DateTime',
|
||||
]
|
||||
timestamp_columns = ([params.date_column] if params.date_column else []) + [
|
||||
c for c in common_timestamp_columns if c != params.date_column
|
||||
]
|
||||
|
||||
for col in timestamp_columns:
|
||||
if col in data.columns:
|
||||
try:
|
||||
data[col] = pd.to_datetime(data[col])
|
||||
data = data.set_index(col)
|
||||
data = data.sort_index()
|
||||
self.logger.info(f'Configured datetime index from column: {col}')
|
||||
return data
|
||||
except (ValueError, TypeError) as e:
|
||||
self.logger.warning(f'Failed to convert column {col} to datetime: {e}')
|
||||
continue
|
||||
|
||||
# If no timestamp column found, check if first column looks like a timestamp
|
||||
first_col = data.columns[0]
|
||||
try:
|
||||
# Try to parse first column as datetime
|
||||
test_values = data[first_col].head(10).dropna()
|
||||
if len(test_values) > 0:
|
||||
pd.to_datetime(test_values)
|
||||
data[first_col] = pd.to_datetime(data[first_col])
|
||||
data = data.set_index(first_col)
|
||||
data = data.sort_index()
|
||||
self.logger.info(f'Configured datetime index from first column: {first_col}')
|
||||
return data
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
|
||||
self.logger.warning('No timestamp column found - some features may not work correctly')
|
||||
return data
|
||||
@@ -40,6 +40,9 @@ with workflow.unsafe.imports_passed_through():
|
||||
from prometheus_client import start_http_server
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import Logger as SientiaLogger
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
from sientia_do.observability.metrics_controller import MetricsController
|
||||
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
||||
|
||||
from model_manager import metrics
|
||||
from model_manager.activities.activities import Activities
|
||||
@@ -48,6 +51,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
build_minio_config,
|
||||
build_mlflow_config,
|
||||
build_mongodb_config,
|
||||
build_plugin_store_config,
|
||||
build_postgres_config,
|
||||
)
|
||||
from model_manager.utils.logger_helper import get_logger
|
||||
@@ -55,6 +59,7 @@ with workflow.unsafe.imports_passed_through():
|
||||
from model_manager.workflows.train_model import TrainModel
|
||||
|
||||
POD_ID = os.getenv('POD_ID')
|
||||
RUNTIME = os.getenv('RUNTIME')
|
||||
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
|
||||
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue')
|
||||
CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue')
|
||||
@@ -79,13 +84,19 @@ async def main():
|
||||
Exception: Any unhandled exception during worker execution
|
||||
SystemExit: On graceful shutdown or error conditions
|
||||
"""
|
||||
|
||||
if not RUNTIME:
|
||||
raise ValueError('RUNTIME environment variable is required')
|
||||
|
||||
host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
|
||||
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
|
||||
logger = get_logger(__name__)
|
||||
|
||||
metadata = {
|
||||
'pod_id': POD_ID,
|
||||
'runtime': RUNTIME,
|
||||
}
|
||||
|
||||
|
||||
start_prometheus_server(logger, metadata)
|
||||
mongo_config = build_mongodb_config()
|
||||
@@ -99,12 +110,41 @@ async def main():
|
||||
|
||||
logger.custom_info(f'MongoDB client initialized at {mongo_config["uri"]}', metadata)
|
||||
|
||||
logger.custom_info(f'Initializing metrics controller', metadata)
|
||||
|
||||
metrics_controller = MetricsController(
|
||||
logger=logger
|
||||
)
|
||||
|
||||
logger.custom_info(f'Installing runtime {RUNTIME}', metadata)
|
||||
|
||||
plugin_store_parameters = build_plugin_store_config()
|
||||
plugin_store = PluginStore(
|
||||
base_url=plugin_store_parameters['base_url'],
|
||||
owner=plugin_store_parameters['owner'],
|
||||
repo=plugin_store_parameters['repo'],
|
||||
username=plugin_store_parameters['username'],
|
||||
password=plugin_store_parameters['password'],
|
||||
branch=plugin_store_parameters['branch'],
|
||||
cache_ttl_seconds=plugin_store_parameters['cache_ttl_seconds'],
|
||||
pypi_index_url=plugin_store_parameters['pypi_index_url'],
|
||||
pypi_username=plugin_store_parameters['pypi_username'],
|
||||
pypi_password=plugin_store_parameters['pypi_password'],
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
await plugin_store.install_runtime(runtime_name=RUNTIME)
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=build_postgres_config(),
|
||||
mlflow_config=build_mlflow_config(),
|
||||
minio_config=build_minio_config(),
|
||||
plugin_store=plugin_store,
|
||||
logger=logger,
|
||||
notification_handler=notification_handler,
|
||||
metrics_controller=metrics_controller,
|
||||
)
|
||||
|
||||
new_runtime = Runtime(
|
||||
@@ -133,44 +173,33 @@ async def main():
|
||||
# The schedule can be created manually if needed
|
||||
|
||||
workers = [
|
||||
Worker(
|
||||
temporal_client,
|
||||
task_queue=TRAIN_TASK_QUEUE,
|
||||
workflows=[TrainModel],
|
||||
prepare_worker(
|
||||
main_workflow=TrainModel,
|
||||
other_workflows=[],
|
||||
activities=[
|
||||
activities.update_experiment_run,
|
||||
activities.validate_train_params,
|
||||
activities.train_model,
|
||||
activities.cleanup_resources,
|
||||
],
|
||||
max_concurrent_workflow_tasks=10,
|
||||
max_concurrent_activities=10,
|
||||
max_concurrent_local_activities=10,
|
||||
max_cached_workflows=100,
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
temporal_client=temporal_client,
|
||||
logger=logger,
|
||||
),
|
||||
Worker(
|
||||
temporal_client,
|
||||
task_queue=CLEANUP_TASK_QUEUE,
|
||||
workflows=[CleanupFiles],
|
||||
prepare_worker(
|
||||
main_workflow=CleanupFiles,
|
||||
other_workflows=[],
|
||||
activities=[
|
||||
activities.cleanup_minio_files,
|
||||
activities.cleanup_temp_directories,
|
||||
],
|
||||
max_concurrent_workflow_tasks=20,
|
||||
max_concurrent_activities=20,
|
||||
max_concurrent_local_activities=20,
|
||||
max_cached_workflows=100,
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
|
||||
temporal_client=temporal_client,
|
||||
logger=logger,
|
||||
),
|
||||
]
|
||||
|
||||
handlers = []
|
||||
|
||||
for w in workers:
|
||||
handlers.append(w.run())
|
||||
handlers = [
|
||||
w.run() for w in workers
|
||||
]
|
||||
|
||||
logger.custom_info('Model manager workers initialized', metadata)
|
||||
|
||||
@@ -183,7 +212,7 @@ async def main():
|
||||
finally:
|
||||
notification_handler.shutdown()
|
||||
logger.custom_info('MongoDB client closed', metadata)
|
||||
await activities.shutdown()
|
||||
activities.shutdown()
|
||||
# Exit with a non-zero status code to indicate failure to Kubernetes
|
||||
metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
|
||||
sys.exit(1)
|
||||
|
||||
@@ -118,7 +118,6 @@ class TrainModel:
|
||||
|
||||
await self._cleanup_resources(
|
||||
experiment_run_id=experiment_run_id,
|
||||
run_dir=(train_result.get('run_dir') or ''),
|
||||
bucket_name=train_params.bucket_name,
|
||||
file_name=train_params.file_name,
|
||||
metadata=metadata,
|
||||
@@ -275,7 +274,6 @@ class TrainModel:
|
||||
async def _cleanup_resources(
|
||||
self,
|
||||
experiment_run_id: int,
|
||||
run_dir: str,
|
||||
bucket_name: str,
|
||||
file_name: str,
|
||||
metadata: dict[str, Any],
|
||||
@@ -283,12 +281,11 @@ class TrainModel:
|
||||
"""
|
||||
Cleanup resources and delete file from MinIO.
|
||||
|
||||
This method removes the temporary run directory via activity and deletes
|
||||
the training file from MinIO. On success, updates DB status to FILE_DELETED.
|
||||
This method deletes the training file from MinIO. On success, updates
|
||||
DB status to FILE_DELETED.
|
||||
On error, updates DB status to FILE_DELETE_ERROR.
|
||||
|
||||
Args:
|
||||
saved_result: TrainModelResult with run_dir and params information
|
||||
experiment_run_id: Validated experiment run ID
|
||||
metadata: Workflow execution metadata
|
||||
|
||||
@@ -300,7 +297,6 @@ class TrainModel:
|
||||
Activities.cleanup_resources,
|
||||
{
|
||||
**metadata,
|
||||
'run_dir': run_dir,
|
||||
'bucket_name': bucket_name,
|
||||
'file_name': file_name,
|
||||
},
|
||||
|
||||
@@ -113,6 +113,10 @@ ignore_missing_imports = true
|
||||
module = "sklearn.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "sientia_model.*"
|
||||
ignore_missing_imports = true
|
||||
|
||||
[[tool.mypy.overrides]]
|
||||
module = "yaml"
|
||||
ignore_missing_imports = true
|
||||
@@ -136,6 +140,9 @@ markers = [
|
||||
"unit: marks tests as unit tests",
|
||||
]
|
||||
|
||||
[tool.pyright]
|
||||
reportMissingTypeStubs = false
|
||||
|
||||
[tool.coverage.run]
|
||||
source = ["model_manager"]
|
||||
omit = [
|
||||
|
||||
@@ -3,10 +3,7 @@ psycopg2-binary==2.9.11
|
||||
sqlalchemy==2.0.44
|
||||
boto3==1.40.55
|
||||
botocore==1.40.55
|
||||
git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1
|
||||
git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.7.2
|
||||
git+https://github.com/Aignosi/sientia-dataops-library.git@v1.10.1
|
||||
prometheus-client==0.23.1
|
||||
mlflow==2.18.0
|
||||
setuptools<81
|
||||
evidently==0.4.21
|
||||
beautifulsoup4==4.12.3
|
||||
scikit-learn==1.4.2
|
||||
|
||||
@@ -289,17 +289,18 @@ def run_local_pipeline(
|
||||
csv_path: Path,
|
||||
save_mlflow: bool = False,
|
||||
) -> dict:
|
||||
"""Run the same training pipeline locally (validate + train + after_train).
|
||||
"""Run the same training pipeline locally (validate + train + metrics).
|
||||
|
||||
Reads CSV from disk, runs TrainingRepository.train and after_train_calculation.
|
||||
Optionally saves to MLflow if save_mlflow is True (requires MLflow env).
|
||||
Reads CSV from disk, runs DataManagerRepository.prepare_training_data and
|
||||
compute_regression_metrics. Optionally saves to MLflow if save_mlflow is
|
||||
True (requires MLflow env).
|
||||
|
||||
Returns:
|
||||
dict: {'success': bool, 'error': str | None, 'scenario': str, ...}
|
||||
"""
|
||||
from model_manager.utils.logger_helper import get_logger
|
||||
from model_manager.utils.models.train_model_params import TrainModelParams
|
||||
from model_manager.utils.repository.training_repository import TrainingRepository
|
||||
from model_manager.utils.repository.data_manager_repository import DataManagerRepository
|
||||
|
||||
result = {
|
||||
'scenario': scenario_name,
|
||||
@@ -327,14 +328,22 @@ def run_local_pipeline(
|
||||
return result
|
||||
|
||||
logger = get_logger(__name__)
|
||||
training_repository = TrainingRepository(logger)
|
||||
data_manager_repository = DataManagerRepository(logger)
|
||||
|
||||
with open(csv_path, 'rb') as f:
|
||||
file_content = BytesIO(f.read())
|
||||
|
||||
try:
|
||||
train_result = training_repository.train(file_content, train_params)
|
||||
train_result = training_repository.after_train_calculation(train_params, train_result)
|
||||
train_result = data_manager_repository.prepare_training_data(
|
||||
train_file_bytes=file_content.getvalue(),
|
||||
validation_file_bytes=None,
|
||||
params=train_params,
|
||||
metadata={'source': 'run_local_pipeline', 'scenario': scenario_name},
|
||||
)
|
||||
train_result = data_manager_repository.compute_regression_metrics(
|
||||
train_params,
|
||||
train_result,
|
||||
)
|
||||
except Exception as e:
|
||||
result['error'] = str(e)
|
||||
raise # re-raise so caller gets full traceback for diagnosis
|
||||
|
||||
@@ -1,316 +0,0 @@
|
||||
"""Unit tests for Activities class with 100% coverage."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_notification_handler():
|
||||
"""Create a mock notification handler."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def postgres_config():
|
||||
"""Create a valid PostgreSQL configuration."""
|
||||
return {
|
||||
'host': 'localhost',
|
||||
'port': 5432,
|
||||
'user': 'testuser',
|
||||
'password': 'testpass',
|
||||
'dbname': 'testdb',
|
||||
'min_connections': 1,
|
||||
'max_connections': 10,
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mlflow_config():
|
||||
"""Create a valid MLFlow configuration."""
|
||||
return {
|
||||
'url': 'http://mlflow:5080',
|
||||
'username': 'aignosi',
|
||||
'password': 'aignosi',
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def minio_config():
|
||||
"""Create a valid MinIO configuration."""
|
||||
return {
|
||||
'endpoint_url': 'http://minio:9000',
|
||||
'access_key': 'minioadmin',
|
||||
'secret_key': 'minioadmin',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'standard',
|
||||
'connect_timeout': 5,
|
||||
'read_timeout': 5,
|
||||
}
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_init_success(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test successful initialization of Activities."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
mock_et_init.assert_called_once()
|
||||
assert mock_et_init.call_args[1]['host'] == postgres_config['host']
|
||||
assert mock_et_init.call_args[1]['port'] == postgres_config['port']
|
||||
assert mock_et_init.call_args[1]['user'] == postgres_config['user']
|
||||
assert mock_et_init.call_args[1]['password'] == postgres_config['password']
|
||||
assert mock_et_init.call_args[1]['dbname'] == postgres_config['dbname']
|
||||
assert mock_et_init.call_args[1]['min_connections'] == postgres_config['min_connections']
|
||||
assert mock_et_init.call_args[1]['max_connections'] == postgres_config['max_connections']
|
||||
assert mock_et_init.call_args[1]['logger'] is mock_logger
|
||||
assert mock_et_init.call_args[1]['notification_handler'] is mock_notification_handler
|
||||
|
||||
mock_model_repo.assert_called_once_with(
|
||||
url=mlflow_config['url'],
|
||||
username=mlflow_config['username'],
|
||||
password=mlflow_config['password'],
|
||||
logger=mock_logger,
|
||||
)
|
||||
|
||||
mock_storage_repo.assert_called_once_with(
|
||||
endpoint_url=minio_config['endpoint_url'],
|
||||
access_key=minio_config['access_key'],
|
||||
secret_key=minio_config['secret_key'],
|
||||
region=minio_config['region'],
|
||||
use_ssl=minio_config['use_ssl'],
|
||||
max_retry_attempts=minio_config['max_retry_attempts'],
|
||||
retry_mode=minio_config['retry_mode'],
|
||||
connect_timeout=minio_config['connect_timeout'],
|
||||
read_timeout=minio_config['read_timeout'],
|
||||
logger=mock_logger,
|
||||
)
|
||||
|
||||
mock_training_init.assert_called_once()
|
||||
assert mock_training_init.call_args[1]['model_repository'] is mock_model_repo.return_value
|
||||
assert mock_training_init.call_args[1]['storage_repository'] is mock_storage_repo.return_value
|
||||
assert mock_training_init.call_args[1]['logger'] is mock_logger
|
||||
assert mock_training_init.call_args[1]['notification_handler'] is mock_notification_handler
|
||||
|
||||
assert hasattr(activities, 'model_repository')
|
||||
assert hasattr(activities, 'storage_repository')
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.close')
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_shutdown(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_close,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test Activities.shutdown() calls ExperimentTracking.close()."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
asyncio.run(activities.shutdown())
|
||||
|
||||
mock_et_close.assert_called_once_with(activities)
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_del_without_engine(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test __del__ when engine attribute does not exist."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
if hasattr(activities, 'engine'):
|
||||
delattr(activities, 'engine')
|
||||
|
||||
activities.__del__()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_del_with_engine_no_super_del(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test __del__ when engine exists but super has no __del__."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
activities.engine = MagicMock()
|
||||
|
||||
with patch('builtins.super') as mock_super:
|
||||
mock_super_instance = MagicMock()
|
||||
del mock_super_instance.__del__
|
||||
mock_super.return_value = mock_super_instance
|
||||
|
||||
activities.__del__()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_del_with_engine_and_super_del(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test __del__ when engine exists and super has __del__."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
activities.engine = MagicMock()
|
||||
|
||||
mock_super_del = MagicMock()
|
||||
|
||||
class MockSuper:
|
||||
def __del__(self):
|
||||
mock_super_del()
|
||||
|
||||
with patch('builtins.super', return_value=MockSuper()):
|
||||
activities.__del__()
|
||||
|
||||
mock_super_del.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.activities.activities.ExperimentTracking.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.Training.__init__', return_value=None)
|
||||
@patch('model_manager.activities.activities.ModelRepository')
|
||||
@patch('model_manager.activities.activities.StorageRepository')
|
||||
def test_activities_del_with_engine_exception_caught(
|
||||
mock_storage_repo,
|
||||
mock_model_repo,
|
||||
mock_training_init,
|
||||
mock_et_init,
|
||||
postgres_config,
|
||||
mlflow_config,
|
||||
minio_config,
|
||||
mock_logger,
|
||||
mock_notification_handler,
|
||||
):
|
||||
"""Test __del__ catches exceptions when super().__del__() raises."""
|
||||
from model_manager.activities.activities import Activities
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=postgres_config,
|
||||
mlflow_config=mlflow_config,
|
||||
minio_config=minio_config,
|
||||
logger=mock_logger,
|
||||
notification_handler=mock_notification_handler,
|
||||
)
|
||||
|
||||
activities.engine = MagicMock()
|
||||
|
||||
class MockSuperWithError:
|
||||
def __del__(self):
|
||||
# Only raise error if not being cleaned up by garbage collector
|
||||
# This prevents the PytestUnraisableExceptionWarning
|
||||
if hasattr(self, '_should_raise') and self._should_raise:
|
||||
raise RuntimeError('Test error')
|
||||
|
||||
# Suppress the PytestUnraisableExceptionWarning for this specific test
|
||||
import warnings
|
||||
|
||||
warnings.filterwarnings('ignore', category=pytest.PytestUnraisableExceptionWarning)
|
||||
|
||||
mock_super = MockSuperWithError()
|
||||
mock_super._should_raise = True
|
||||
try:
|
||||
with patch('builtins.super', return_value=mock_super):
|
||||
activities.__del__()
|
||||
finally:
|
||||
# Prevent the exception from being raised during garbage collection
|
||||
mock_super._should_raise = False
|
||||
@@ -61,7 +61,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
|
||||
x_train, x_test, y_train, y_test = sample_dataframes
|
||||
process_data = MagicMock()
|
||||
regr = MagicMock()
|
||||
scaler_dict = {'var1': {'min': 0, 'max': 100}}
|
||||
|
||||
result = TrainModelResult(
|
||||
params=sample_params,
|
||||
@@ -71,7 +70,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=regr,
|
||||
scaler_dict=scaler_dict,
|
||||
)
|
||||
|
||||
assert result.params == sample_params
|
||||
@@ -96,7 +94,6 @@ def test_train_model_result_optional_fields_default_none(sample_params, sample_d
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=MagicMock(),
|
||||
scaler_dict={},
|
||||
)
|
||||
|
||||
assert result.y_pred is None
|
||||
@@ -123,7 +120,6 @@ def test_train_model_result_with_metrics(sample_params, sample_dataframes):
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=MagicMock(),
|
||||
scaler_dict={},
|
||||
y_pred=y_pred,
|
||||
mse_val=1.5,
|
||||
mae_val=1.2,
|
||||
@@ -148,7 +144,6 @@ def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=MagicMock(),
|
||||
scaler_dict={},
|
||||
run_name='test-experiment-1',
|
||||
report_path='/path/to/report.html',
|
||||
train_data_path='/path/to/train_data.csv',
|
||||
@@ -186,11 +181,11 @@ def test_train_model_result_is_dataclass(sample_params, sample_dataframes):
|
||||
|
||||
|
||||
def test_train_model_result_field_count():
|
||||
"""Test that TrainModelResult has exactly 20 fields."""
|
||||
"""Test that TrainModelResult has exactly 19 fields."""
|
||||
from dataclasses import fields
|
||||
|
||||
result_fields = fields(TrainModelResult)
|
||||
assert len(result_fields) == 20
|
||||
assert len(result_fields) == 19
|
||||
|
||||
field_names = {f.name for f in result_fields}
|
||||
expected_fields = {
|
||||
@@ -201,7 +196,6 @@ def test_train_model_result_field_count():
|
||||
'y_train',
|
||||
'y_test',
|
||||
'regr',
|
||||
'scaler_dict',
|
||||
'y_pred',
|
||||
'y_train_pred',
|
||||
'mse_val',
|
||||
@@ -231,7 +225,6 @@ def test_train_model_result_complete_workflow(sample_params, sample_dataframes):
|
||||
y_train=y_train,
|
||||
y_test=y_test,
|
||||
regr=MagicMock(),
|
||||
scaler_dict={'var1': {'min': 0, 'max': 100}},
|
||||
)
|
||||
|
||||
# Step 2: Add predictions and metrics
|
||||
|
||||
@@ -1,982 +0,0 @@
|
||||
"""Unit tests for ModelRepository with 100% coverage."""
|
||||
|
||||
import os
|
||||
import shutil
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_temp_directories():
|
||||
"""Clean up temporary directories after each test."""
|
||||
# Get the temp directory path
|
||||
current_file_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
model_manager_dir = os.path.dirname(os.path.dirname(os.path.dirname(current_file_dir)))
|
||||
temp_dir = os.path.join(model_manager_dir, 'reports', 'temp')
|
||||
|
||||
# Run the test
|
||||
yield
|
||||
|
||||
# Clean up after test
|
||||
if os.path.exists(temp_dir):
|
||||
for item in os.listdir(temp_dir):
|
||||
item_path = os.path.join(temp_dir, item)
|
||||
if os.path.isdir(item_path) and item.startswith('test_run_'):
|
||||
try:
|
||||
shutil.rmtree(item_path)
|
||||
except (OSError, PermissionError):
|
||||
# Ignore cleanup errors
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger."""
|
||||
return MagicMock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_train_result():
|
||||
"""Create a mock TrainModelResult."""
|
||||
result = MagicMock()
|
||||
result.params = MagicMock()
|
||||
result.params.experiment_name = 'test_experiment'
|
||||
result.params.experiment_run_id = 1
|
||||
result.params.target_variable = 'target'
|
||||
result.params.variable_columns = ['var1', 'var2']
|
||||
result.params.lag_train = 5
|
||||
result.params.lag_val = 3
|
||||
result.params.window = 10
|
||||
result.params.low_lim = {'var1': 0.0}
|
||||
result.params.upp_lim = {'var1': 10.0}
|
||||
result.params.include_ar = False
|
||||
result.params.train_size = 80
|
||||
result.params.removed_intervals = []
|
||||
result.params.rem_static_win = True
|
||||
result.params.static_threshold = None
|
||||
result.run_name = 'test_run'
|
||||
result.run_dir = '/tmp/test_run' # noqa: S108
|
||||
result.report_path = '/tmp/test_run/report.html' # noqa: S108
|
||||
result.train_data_path = '/tmp/test_run/train_data.csv' # noqa: S108
|
||||
result.test_data_path = '/tmp/test_run/test_data.csv' # noqa: S108
|
||||
result.mse_val = 0.5
|
||||
result.r2_val = 0.9
|
||||
result.mae_val = 0.3
|
||||
result.scaler_dict = {'scaler': 'standard'}
|
||||
result.process_data = MagicMock()
|
||||
result.regr = MagicMock()
|
||||
result.regr.predict = MagicMock(return_value=np.array([1.0, 2.0, 3.0]))
|
||||
result.x_train = pd.DataFrame({'var1': [1, 2, 3]})
|
||||
result.y_train = pd.Series([1.0, 2.0, 3.0], name='target')
|
||||
result.x_test = pd.DataFrame({'var1': [4, 5, 6]})
|
||||
result.y_test = pd.Series([4.0, 5.0, 6.0], name='target')
|
||||
result.y_pred = np.array([4.1, 5.1, 6.1])
|
||||
return result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_model_repository_init(mock_model_serving_class, mock_logger):
|
||||
"""Test ModelRepository initialization."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_model_serving_class.assert_called_once_with(
|
||||
tracking_uri='http://mlflow.test', username='user', password='pass'
|
||||
)
|
||||
assert repo.model_serving is mock_model_serving_instance
|
||||
assert repo.logger is mock_logger
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_save_model_success(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test save_model successfully saves model."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
repo._get_next_run_name = MagicMock(return_value='test_experiment-1')
|
||||
repo._generate_artifacts = MagicMock(return_value=mock_train_result)
|
||||
repo._save_run = MagicMock()
|
||||
|
||||
result = repo.save_model(mock_train_result)
|
||||
|
||||
repo._get_next_run_name.assert_called_once_with('test_experiment')
|
||||
repo._generate_artifacts.assert_called_once()
|
||||
repo._save_run.assert_called_once_with(mock_train_result)
|
||||
mock_logger.info.assert_called_once()
|
||||
assert result is mock_train_result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.os.path.exists')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.rmtree')
|
||||
def test_cleanup_run_directory_exists(
|
||||
mock_rmtree, mock_exists, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test cleanup_run_directory when directory exists."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
||||
|
||||
mock_exists.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_rmtree.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_logger.info.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.os.path.exists')
|
||||
def test_cleanup_run_directory_not_exists(mock_exists, mock_model_serving_class, mock_logger):
|
||||
"""Test cleanup_run_directory when directory doesn't exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
mock_exists.return_value = False
|
||||
|
||||
repo.cleanup_run_directory('/tmp/test_run') # noqa: S108
|
||||
|
||||
mock_exists.assert_called_once_with('/tmp/test_run') # noqa: S108
|
||||
mock_logger.info.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_cleanup_run_directory_empty_path(mock_model_serving_class, mock_logger):
|
||||
"""Test cleanup_run_directory with empty path."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Reset mock after initialization to focus on method-specific calls
|
||||
mock_logger.reset_mock()
|
||||
|
||||
repo.cleanup_run_directory('')
|
||||
|
||||
mock_logger.info.assert_called_once_with('No run directory specified, skipping cleanup')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_next_run_name_no_existing_runs(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_next_run_name when no runs exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_instance.search_runs_by_name.return_value = []
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_next_run_name('test_experiment')
|
||||
|
||||
assert result == 'test_experiment-1'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_next_run_name_with_existing_runs(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_next_run_name when runs exist."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_instance.search_runs_by_name.return_value = [
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
MagicMock(),
|
||||
]
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_next_run_name('test_experiment')
|
||||
|
||||
assert result == 'test_experiment-4'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_get_reports_directory(mock_model_serving_class, mock_logger):
|
||||
"""Test _get_reports_directory returns correct path."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
result = repo._get_reports_directory()
|
||||
|
||||
assert result.endswith(os.path.join('model_manager', 'reports'))
|
||||
assert os.path.isabs(result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_success(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data successfully prepares data."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data, current_data = repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
# Check reference data
|
||||
assert 'target' in reference_data.columns
|
||||
assert 'prediction' in reference_data.columns
|
||||
assert len(reference_data) == 3
|
||||
|
||||
# Check current data
|
||||
assert 'target' in current_data.columns
|
||||
assert 'prediction' in current_data.columns
|
||||
assert len(current_data) == 3
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_x_train(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when x_train is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.x_train = pd.DataFrame()
|
||||
|
||||
with pytest.raises(ValueError, match='Training features .* are empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_y_train(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when y_train is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_train = pd.Series(dtype=float)
|
||||
|
||||
with pytest.raises(ValueError, match='Training target .* is empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_none_y_pred(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when y_pred is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_pred = None
|
||||
|
||||
with pytest.raises(ValueError, match='Test predictions .* are None'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_none_y_train_pred(
|
||||
mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _init_artifacts_data raises ValueError when y_train_pred is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_train_pred = None
|
||||
|
||||
with pytest.raises(ValueError, match='Training predictions .* are None'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.datetime')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_success(
|
||||
mock_makedirs, mock_datetime, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _create_run_directory creates directory successfully."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_datetime.now.return_value.strftime.return_value = '20240101_120000_123456'
|
||||
|
||||
result = repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
expected_path = os.path.normpath(
|
||||
os.path.join('/tmp/reports', 'temp', 'test_run_20240101_120000_123456') # noqa: S108
|
||||
)
|
||||
assert os.path.normpath(result) == expected_path
|
||||
mock_makedirs.assert_called_once()
|
||||
call_path = mock_makedirs.call_args[0][0]
|
||||
assert os.path.normpath(call_path) == expected_path
|
||||
assert mock_makedirs.call_args[1] == {'exist_ok': True}
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_permission_error(
|
||||
mock_makedirs, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _create_run_directory raises PermissionError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_makedirs.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when creating directory'):
|
||||
repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_success(mock_open, mock_copy, mock_model_serving_class, mock_logger):
|
||||
"""Test _setup_run_directory creates files successfully."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
# Check that empty files were created
|
||||
assert mock_open.call_count == 3
|
||||
mock_copy.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.json.dump')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_with_equation(
|
||||
mock_open,
|
||||
mock_reports_class,
|
||||
mock_json_dump,
|
||||
mock_model_serving_class,
|
||||
mock_logger,
|
||||
mock_train_result,
|
||||
):
|
||||
"""Test _generate_report creates equation JSON artifact."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Add equation to train result
|
||||
mock_train_result.equation = {
|
||||
'target_variable': 'target',
|
||||
'coefficients': {'var1': 1.5, 'var2': -0.75},
|
||||
'intercept': 10.5,
|
||||
'equation_string': 'target = 10.5 + 1.5 * var1 + -0.75 * var2',
|
||||
'latex_equation': 'target = 10.5 + 1.5 \\cdot var1 + -0.75 \\cdot var2',
|
||||
'model_type': 'Linear Regression',
|
||||
}
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock DataFrame.to_csv to avoid file I/O
|
||||
with patch.object(pd.DataFrame, 'to_csv'):
|
||||
result = repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
# Verify equation path was set
|
||||
assert result.equation_path == os.path.join(mock_train_result.run_dir, 'model_equation.json')
|
||||
|
||||
# Verify JSON was written
|
||||
mock_json_dump.assert_called()
|
||||
call_args = mock_json_dump.call_args
|
||||
assert call_args[0][0] == mock_train_result.equation
|
||||
assert call_args[1]['indent'] == 2
|
||||
assert call_args[1]['ensure_ascii'] is False
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_without_equation(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report works without equation."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# No equation
|
||||
mock_train_result.equation = None
|
||||
# Remove equation_path if it exists from fixture
|
||||
if hasattr(mock_train_result, 'equation_path'):
|
||||
del mock_train_result.equation_path
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock DataFrame.to_csv to avoid file I/O
|
||||
with patch.object(pd.DataFrame, 'to_csv'):
|
||||
result = repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
# Verify equation section was not executed (equation_path not set)
|
||||
# Since equation is None, the equation block should not run
|
||||
assert result == mock_train_result
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_equation(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs equation artifact."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set equation path
|
||||
mock_train_result.equation_path = '/tmp/test_run/model_equation.json' # noqa: S108
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify equation artifact was logged
|
||||
logged_artifacts = [
|
||||
call[0][0] for call in mock_model_serving_instance.log_artifact.call_args_list
|
||||
]
|
||||
assert '/tmp/test_run/model_equation.json' in logged_artifacts # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_static_threshold_value(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs static_threshold when rem_static_win is True and value is set."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set static_threshold to a specific value
|
||||
mock_train_result.params.rem_static_win = True
|
||||
mock_train_result.params.static_threshold = 500
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify static_threshold was logged with the correct value
|
||||
log_param_calls = {
|
||||
call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list
|
||||
}
|
||||
assert log_param_calls['static_threshold'] == 500
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_with_rem_static_win_false(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run logs static_threshold as None when rem_static_win is False."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Set rem_static_win to False
|
||||
mock_train_result.params.rem_static_win = False
|
||||
mock_train_result.params.static_threshold = 500 # Should be ignored
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock all path.exists calls to return True
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify static_threshold was logged as None
|
||||
log_param_calls = {
|
||||
call[0][0]: call[0][1] for call in mock_model_serving_instance.log_param.call_args_list
|
||||
}
|
||||
assert log_param_calls['static_threshold'] is None
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_without_equation(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run works without equation."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
mock_model_serving_instance = MagicMock()
|
||||
mock_model_serving_class.return_value = mock_model_serving_instance
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# No equation
|
||||
mock_train_result.equation_path = None
|
||||
|
||||
# Mock path.exists to return True for required artifacts
|
||||
mock_exists.return_value = True
|
||||
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
# Verify only 3 artifacts were logged (report, train_data, test_data)
|
||||
assert mock_model_serving_instance.log_artifact.call_count == 3
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_report(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when report is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock report doesn't exist
|
||||
def exists_side_effect(path):
|
||||
return not path.endswith('report.html')
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Report file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_none_metrics(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when metrics are None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.mse_val = None
|
||||
# Mock all paths exist so we reach the metrics check
|
||||
mock_exists.return_value = True
|
||||
|
||||
with pytest.raises(ValueError, match='One or more metrics .* are None'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_train_data(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when train data is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock train_data doesn't exist
|
||||
def exists_side_effect(path):
|
||||
return not path.endswith('train_data.csv')
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Training data file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_save_run_missing_test_data(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _save_run raises ValueError when test data is missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock test_data doesn't exist
|
||||
def exists_side_effect(path):
|
||||
if path.endswith('test_data.csv'):
|
||||
return False
|
||||
return True
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
with pytest.raises(ValueError, match='Test data file does not exist'):
|
||||
repo._save_run(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_x_test(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when x_test is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.x_test = pd.DataFrame()
|
||||
|
||||
with pytest.raises(ValueError, match='Test features .* are empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
def test_init_artifacts_data_empty_y_test(mock_model_serving_class, mock_logger, mock_train_result):
|
||||
"""Test _init_artifacts_data raises ValueError when y_test is empty."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.y_test = pd.Series(dtype=float)
|
||||
|
||||
with pytest.raises(ValueError, match='Test target .* is empty'):
|
||||
repo._init_artifacts_data(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.makedirs')
|
||||
def test_create_run_directory_os_error(mock_makedirs, mock_model_serving_class, mock_logger):
|
||||
"""Test _create_run_directory raises OSError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_makedirs.side_effect = OSError('Disk full')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to create directory'):
|
||||
repo._create_run_directory('/tmp/reports', 'test_run') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_file_not_found(
|
||||
mock_open, mock_copy, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _setup_run_directory raises FileNotFoundError when header missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_copy.side_effect = FileNotFoundError('Header not found')
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Header file not found'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_permission_error(
|
||||
mock_open, mock_copy, mock_model_serving_class, mock_logger
|
||||
):
|
||||
"""Test _setup_run_directory raises PermissionError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_open.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when setting up directory'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.shutil.copy')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_setup_run_directory_os_error(mock_open, mock_copy, mock_model_serving_class, mock_logger):
|
||||
"""Test _setup_run_directory raises OSError."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_open.side_effect = OSError('Disk error')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to setup run directory'):
|
||||
repo._setup_run_directory('/tmp/test_run', '/tmp/header.html') # noqa: S108
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_value_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises ValueError on invalid data."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Create data that can't be converted to float64
|
||||
reference_data = pd.DataFrame({'var1': ['invalid', 'data']})
|
||||
current_data = pd.DataFrame({'var1': [1, 2]})
|
||||
|
||||
with pytest.raises(ValueError, match='Failed to convert data to float64'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_permission_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises PermissionError on write failure."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock Reports to raise PermissionError
|
||||
mock_reports_class.side_effect = PermissionError('Permission denied')
|
||||
|
||||
with pytest.raises(PermissionError, match='Permission denied when writing report files'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_os_error(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises OSError on write failure."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
# Mock Reports to raise OSError
|
||||
mock_reports_class.side_effect = OSError('Disk error')
|
||||
|
||||
with pytest.raises(OSError, match='Failed to generate report'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.Reports')
|
||||
@patch('builtins.open', create=True)
|
||||
def test_generate_report_none_run_dir(
|
||||
mock_open, mock_reports_class, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_report raises ValueError when run_dir is None."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.run_dir = None
|
||||
|
||||
reference_data = pd.DataFrame({'var1': [1, 2], 'var2': [3, 4], 'target': [5, 6]})
|
||||
current_data = pd.DataFrame({'var1': [7, 8], 'var2': [9, 10], 'target': [11, 12]})
|
||||
|
||||
with pytest.raises(ValueError, match='run_dir is not set'):
|
||||
repo._generate_report(reference_data, current_data, mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_no_run_name(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises ValueError when run_name is not set."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
mock_train_result.run_name = None
|
||||
mock_exists.return_value = True # Mock reports directory exists
|
||||
|
||||
# Mock _create_run_directory to avoid creating real directories
|
||||
with patch.object(repo, '_create_run_directory') as mock_create_dir:
|
||||
mock_create_dir.return_value = '/mock/run/dir'
|
||||
|
||||
with pytest.raises(ValueError, match='run_name must be set'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_reports_dir_not_found(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises FileNotFoundError when reports dir missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock reports directory doesn't exist
|
||||
mock_exists.return_value = False
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Reports directory does not exist'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
def test_generate_artifacts_header_not_found(
|
||||
mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts raises FileNotFoundError when header missing."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock: reports dir exists, but header doesn't
|
||||
def exists_side_effect(path):
|
||||
if path.endswith('header.html'):
|
||||
return False
|
||||
return True
|
||||
|
||||
mock_exists.side_effect = exists_side_effect
|
||||
|
||||
# Mock _create_run_directory to avoid creating real directories
|
||||
with patch.object(repo, '_create_run_directory') as mock_create_dir:
|
||||
mock_create_dir.return_value = '/mock/run/dir'
|
||||
|
||||
with pytest.raises(FileNotFoundError, match='Header file does not exist'):
|
||||
repo._generate_artifacts(mock_train_result)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.model_repository.ModelServing')
|
||||
@patch('model_manager.utils.repository.model_repository.path.exists')
|
||||
@patch('model_manager.utils.repository.model_repository.path.join')
|
||||
def test_generate_artifacts_success(
|
||||
mock_join, mock_exists, mock_model_serving_class, mock_logger, mock_train_result
|
||||
):
|
||||
"""Test _generate_artifacts success case covering lines 147-148."""
|
||||
from model_manager.utils.repository.model_repository import ModelRepository
|
||||
|
||||
repo = ModelRepository(
|
||||
url='http://mlflow.test', username='user', password='pass', logger=mock_logger
|
||||
)
|
||||
|
||||
# Mock path.join to return predictable paths
|
||||
def join_side_effect(*args):
|
||||
return '/'.join(args)
|
||||
|
||||
mock_join.side_effect = join_side_effect
|
||||
mock_exists.return_value = True # Both reports dir and header.html exist
|
||||
|
||||
# Mock the internal methods to avoid actual file operations
|
||||
with (
|
||||
patch.object(repo, '_setup_run_directory') as mock_setup,
|
||||
patch.object(repo, '_generate_report') as mock_generate_report,
|
||||
patch.object(repo, '_init_artifacts_data') as mock_init_data,
|
||||
patch.object(repo, '_get_reports_directory') as mock_get_reports_dir,
|
||||
patch.object(repo, '_create_run_directory') as mock_create_run_dir,
|
||||
):
|
||||
# Setup mocks
|
||||
mock_init_data.return_value = (pd.DataFrame(), pd.DataFrame())
|
||||
mock_get_reports_dir.return_value = '/reports'
|
||||
mock_create_run_dir.return_value = '/reports/run_1'
|
||||
mock_generate_report.return_value = mock_train_result
|
||||
|
||||
# Call the method
|
||||
result = repo._generate_artifacts(mock_train_result)
|
||||
|
||||
# Verify the methods on lines 147-148 were called
|
||||
mock_setup.assert_called_once_with('/reports/run_1', '/reports/header.html')
|
||||
mock_generate_report.assert_called_once()
|
||||
|
||||
# Verify result
|
||||
assert result == mock_train_result
|
||||
@@ -1,513 +0,0 @@
|
||||
"""Unit tests for StorageRepository class."""
|
||||
|
||||
from io import BytesIO
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from botocore.exceptions import ClientError
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_logger():
|
||||
"""Create a mock logger for testing."""
|
||||
logger = Mock()
|
||||
logger.info = Mock()
|
||||
logger.error = Mock()
|
||||
logger.warning = Mock()
|
||||
return logger
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def storage_config():
|
||||
"""Create storage repository configuration."""
|
||||
return {
|
||||
'endpoint_url': 'http://localhost:9000',
|
||||
'access_key': 'test_access_key',
|
||||
'secret_key': 'test_secret_key',
|
||||
'region': 'us-east-1',
|
||||
'use_ssl': False,
|
||||
'max_retry_attempts': 3,
|
||||
'retry_mode': 'standard',
|
||||
'connect_timeout': 30,
|
||||
'read_timeout': 60,
|
||||
}
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_initialization(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository initialization with correct parameters."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Verify attributes are set correctly
|
||||
assert repo.endpoint_url == storage_config['endpoint_url']
|
||||
assert repo.access_key == storage_config['access_key']
|
||||
assert repo.secret_key == storage_config['secret_key']
|
||||
assert repo.region == storage_config['region']
|
||||
assert repo.use_ssl == storage_config['use_ssl']
|
||||
assert repo.max_retry_attempts == storage_config['max_retry_attempts']
|
||||
assert repo.retry_mode == storage_config['retry_mode']
|
||||
assert repo.connect_timeout == storage_config['connect_timeout']
|
||||
assert repo.read_timeout == storage_config['read_timeout']
|
||||
assert repo.logger == mock_logger
|
||||
|
||||
# Verify boto3 client was created
|
||||
mock_boto3.client.assert_called_once()
|
||||
call_args = mock_boto3.client.call_args
|
||||
|
||||
assert call_args[0][0] == 's3'
|
||||
assert call_args[1]['endpoint_url'] == storage_config['endpoint_url']
|
||||
assert call_args[1]['aws_access_key_id'] == storage_config['access_key']
|
||||
assert call_args[1]['aws_secret_access_key'] == storage_config['secret_key']
|
||||
assert call_args[1]['use_ssl'] == storage_config['use_ssl']
|
||||
|
||||
# Verify logger was called
|
||||
mock_logger.info.assert_called()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_boto_config(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that boto3 Config is created with correct retry settings."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Verify Config was passed with correct settings
|
||||
call_args = mock_boto3.client.call_args
|
||||
boto_config = call_args[1]['config']
|
||||
|
||||
assert boto_config.region_name == storage_config['region']
|
||||
assert boto_config.connect_timeout == storage_config['connect_timeout']
|
||||
assert boto_config.read_timeout == storage_config['read_timeout']
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_success(mock_boto3, mock_logger, storage_config):
|
||||
"""Test successful file fetch from MinIO."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response
|
||||
file_content = b'test file content'
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert result.getvalue() == file_content
|
||||
|
||||
# Verify get_object was called correctly
|
||||
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
|
||||
|
||||
# Verify logging
|
||||
assert mock_logger.info.call_count >= 2 # Init + fetch
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_with_large_content(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with large content."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response with large content
|
||||
large_content = b'x' * 1024 * 1024 # 1MB
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = large_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'large-file.bin')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert len(result.getvalue()) == 1024 * 1024
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_empty_content(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with empty content."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock response with empty content
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = b''
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file
|
||||
result = repo.fetch_file('test-bucket', 'empty-file.txt')
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, BytesIO)
|
||||
assert result.getvalue() == b''
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_not_found(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when object doesn't exist."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise NoSuchKey error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'NoSuchKey', 'Message': 'The specified key does not exist.'}},
|
||||
'GetObject',
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch non-existent file
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.fetch_file('test-bucket', 'non-existent.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'NoSuchKey'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_access_denied(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when access is denied."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise AccessDenied error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'GetObject'
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch file without permissions
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.fetch_file('test-bucket', 'protected-file.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file when network error occurs."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise network error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.side_effect = ConnectionError('Network unreachable')
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to fetch file with network error
|
||||
with pytest.raises(ConnectionError):
|
||||
repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_success(mock_boto3, mock_logger, storage_config):
|
||||
"""Test successful file deletion from MinIO."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify delete_object was called correctly
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
|
||||
|
||||
# Verify logging
|
||||
assert mock_logger.info.call_count >= 2 # Init + delete
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_non_existent(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file that doesn't exist (should succeed silently in S3)."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# S3/MinIO delete is idempotent - deleting non-existent file succeeds
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete non-existent file (should succeed)
|
||||
repo.delete_file('test-bucket', 'non-existent.csv')
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once()
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_access_denied(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when access is denied."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise AccessDenied error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ClientError(
|
||||
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'DeleteObject'
|
||||
)
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file without permissions
|
||||
with pytest.raises(ClientError) as exc_info:
|
||||
repo.delete_file('test-bucket', 'protected-file.csv')
|
||||
|
||||
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_network_error(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file when network error occurs."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
# Setup mock to raise network error
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.side_effect = ConnectionError('Network unreachable')
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Attempt to delete file with network error
|
||||
with pytest.raises(ConnectionError):
|
||||
repo.delete_file('test-bucket', 'test-file.csv')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_with_ssl(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository initialization with SSL enabled."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['use_ssl'] = True
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.use_ssl is True
|
||||
|
||||
# Verify boto3 client was created with use_ssl=True
|
||||
call_args = mock_boto3.client.call_args
|
||||
assert call_args[1]['use_ssl'] is True
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_timeouts(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with custom timeout values."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['connect_timeout'] = 10
|
||||
storage_config['read_timeout'] = 120
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.connect_timeout == 10
|
||||
assert repo.read_timeout == 120
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_retry_mode(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different retry modes."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
for retry_mode in ['standard', 'legacy', 'adaptive']:
|
||||
storage_config['retry_mode'] = retry_mode
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
assert repo.retry_mode == retry_mode
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_custom_max_retries(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different max retry attempts."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
storage_config['max_retry_attempts'] = 5
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
assert repo.max_retry_attempts == 5
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_with_special_characters(mock_boto3, mock_logger, storage_config):
|
||||
"""Test fetch file with special characters in name."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
file_content = b'test content'
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Fetch file with special characters
|
||||
special_filename = 'test file (2023-01-01) #1.csv'
|
||||
result = repo.fetch_file('test-bucket', special_filename)
|
||||
|
||||
assert isinstance(result, BytesIO)
|
||||
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key=special_filename)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_delete_file_with_path_separators(mock_boto3, mock_logger, storage_config):
|
||||
"""Test delete file with path separators in object key."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.delete_object.return_value = {}
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
# Delete file with path separators
|
||||
file_path = 'data/2023/01/test-file.csv'
|
||||
repo.delete_file('test-bucket', file_path)
|
||||
|
||||
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key=file_path)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_storage_repository_different_regions(mock_boto3, mock_logger, storage_config):
|
||||
"""Test StorageRepository with different AWS regions."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
regions = ['us-west-1', 'eu-central-1', 'ap-southeast-1']
|
||||
|
||||
for region in regions:
|
||||
storage_config['region'] = region
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
assert repo.region == region
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_fetch_file_logs_file_size(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that fetch_file logs the file size."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
file_content = b'x' * 12345
|
||||
mock_body = Mock()
|
||||
mock_body.read.return_value = file_content
|
||||
mock_body.__enter__ = Mock(return_value=mock_body)
|
||||
mock_body.__exit__ = Mock(return_value=False)
|
||||
|
||||
mock_response = {'Body': mock_body}
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_s3_client.get_object.return_value = mock_response
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
|
||||
repo.fetch_file('test-bucket', 'test-file.csv')
|
||||
|
||||
# Verify logging includes file size
|
||||
log_calls = [str(call) for call in mock_logger.info.call_args_list]
|
||||
assert any('12345 bytes' in str(call) for call in log_calls)
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_close_method(mock_boto3, mock_logger, storage_config):
|
||||
"""Test that the close method calls the underlying client's close method."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
repo.close()
|
||||
|
||||
mock_s3_client.close.assert_called_once()
|
||||
mock_logger.info.assert_called_with('MinIO client closed')
|
||||
|
||||
|
||||
@patch('model_manager.utils.repository.storage_repository.boto3')
|
||||
def test_list_bucket_objects_with_pagination(mock_boto3, mock_logger, storage_config):
|
||||
"""Test list_bucket_objects with a paginated response."""
|
||||
from model_manager.utils.repository.storage_repository import StorageRepository
|
||||
|
||||
mock_s3_client = Mock()
|
||||
mock_paginator = Mock()
|
||||
page1 = {
|
||||
'Contents': [
|
||||
{'Key': 'file1.txt'},
|
||||
{'Key': 'file2.txt'},
|
||||
]
|
||||
}
|
||||
page2 = {
|
||||
'Contents': [
|
||||
{'Key': 'file3.txt'},
|
||||
]
|
||||
}
|
||||
page3 = {}
|
||||
|
||||
mock_paginator.paginate.return_value = [page1, page2, page3]
|
||||
mock_s3_client.get_paginator.return_value = mock_paginator
|
||||
mock_boto3.client.return_value = mock_s3_client
|
||||
|
||||
repo = StorageRepository(logger=mock_logger, **storage_config)
|
||||
objects = repo.list_bucket_objects('test-bucket', max_keys=2)
|
||||
|
||||
assert objects == ['file1.txt', 'file2.txt', 'file3.txt']
|
||||
assert len(objects) == 3
|
||||
mock_s3_client.get_paginator.assert_called_once_with('list_objects_v2')
|
||||
mock_paginator.paginate.assert_called_once_with(Bucket='test-bucket', MaxKeys=2)
|
||||
mock_logger.info.assert_any_call('Listed 3 objects from bucket test-bucket')
|
||||
File diff suppressed because it is too large
Load Diff
@@ -20,6 +20,10 @@ def mock_env_vars():
|
||||
'PROJECT_NAME': 'test-project',
|
||||
'TRAIN_TASK_QUEUE': 'train_model-local_queue',
|
||||
'CLEANUP_TASK_QUEUE': 'cleanup-local_queue',
|
||||
'RUNTIME': 'model-manager-worker',
|
||||
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'STORE_OWNER': 'sientia',
|
||||
'STORE_REPO': 'model-library-store',
|
||||
}
|
||||
|
||||
with patch.dict(os.environ, env_vars, clear=False):
|
||||
@@ -176,6 +180,8 @@ def test_start_prometheus_server_failure(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -191,6 +197,8 @@ async def test_main_successful_startup(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -220,6 +228,23 @@ async def test_main_successful_startup(
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
mock_activities_class.return_value = mock_activities
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
mock_runtime = Mock()
|
||||
mock_runtime_class.return_value = mock_runtime
|
||||
|
||||
@@ -262,6 +287,8 @@ async def test_main_successful_startup(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -277,6 +304,8 @@ async def test_main_handles_exception(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -320,6 +349,23 @@ async def test_main_handles_exception(
|
||||
mock_app_up = Mock()
|
||||
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
# Run main() and expect SystemExit
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
await main()
|
||||
@@ -343,6 +389,8 @@ async def test_main_handles_exception(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -358,6 +406,8 @@ async def test_main_temporal_client_configuration(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -377,6 +427,10 @@ async def test_main_temporal_client_configuration(
|
||||
'TEMPORAL_HOST': 'temporal.example.com:7233',
|
||||
'TEMPORAL_NAMESPACE': 'production',
|
||||
'TEMPORAL_USE_TLS': 'true',
|
||||
'RUNTIME': 'model-manager-worker',
|
||||
'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'STORE_OWNER': 'sientia',
|
||||
'STORE_REPO': 'model-library-store',
|
||||
},
|
||||
):
|
||||
# Setup mocks
|
||||
@@ -390,6 +444,23 @@ async def test_main_temporal_client_configuration(
|
||||
mock_build_mlflow.return_value = {}
|
||||
mock_build_minio.return_value = {}
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
mock_notification_handler = Mock()
|
||||
mock_notification_handler.shutdown = Mock()
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
@@ -431,6 +502,8 @@ async def test_main_temporal_client_configuration(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -446,6 +519,8 @@ async def test_main_worker_configuration(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -500,6 +575,23 @@ async def test_main_worker_configuration(
|
||||
mock_app_up = Mock()
|
||||
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
# Run main()
|
||||
with pytest.raises(SystemExit):
|
||||
await main()
|
||||
@@ -539,6 +631,8 @@ async def test_main_worker_configuration(
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@@ -554,6 +648,8 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
@@ -603,6 +699,106 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
|
||||
mock_app_up = Mock()
|
||||
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch('model_manager.worker.worker.Worker')
|
||||
@patch('model_manager.worker.worker.client.Client')
|
||||
@patch('model_manager.worker.worker.Runtime')
|
||||
@patch('model_manager.worker.worker.Activities')
|
||||
@patch('model_manager.worker.worker.NotificationHandler')
|
||||
@patch('model_manager.worker.worker.PluginStore')
|
||||
@patch('model_manager.worker.worker.build_plugin_store_config')
|
||||
@patch('model_manager.worker.worker.build_mongodb_config')
|
||||
@patch('model_manager.worker.worker.build_postgres_config')
|
||||
@patch('model_manager.worker.worker.build_mlflow_config')
|
||||
@patch('model_manager.worker.worker.build_minio_config')
|
||||
@patch('model_manager.worker.worker.get_logger')
|
||||
@patch('model_manager.worker.worker.start_prometheus_server')
|
||||
@patch('model_manager.worker.worker.metrics')
|
||||
async def test_main_missing_runtime_fails_fast(
|
||||
mock_metrics,
|
||||
mock_start_prometheus,
|
||||
mock_get_logger,
|
||||
mock_build_minio,
|
||||
mock_build_mlflow,
|
||||
mock_build_postgres,
|
||||
mock_build_mongodb,
|
||||
mock_build_plugin_store_config,
|
||||
mock_plugin_store_class,
|
||||
mock_notification_handler_class,
|
||||
mock_activities_class,
|
||||
mock_runtime_class,
|
||||
mock_client_class,
|
||||
mock_worker_class,
|
||||
mock_logger,
|
||||
):
|
||||
"""Test that main() fails fast when RUNTIME is missing."""
|
||||
from model_manager.worker.worker import main
|
||||
|
||||
mock_get_logger.return_value = mock_logger
|
||||
mock_build_mongodb.return_value = {
|
||||
'connection_string': 'mongodb://test',
|
||||
'database_name': 'test_db',
|
||||
'uri': 'localhost:27018',
|
||||
}
|
||||
mock_build_postgres.return_value = {}
|
||||
mock_build_mlflow.return_value = {}
|
||||
mock_build_minio.return_value = {}
|
||||
|
||||
mock_notification_handler = Mock()
|
||||
mock_notification_handler.shutdown = Mock()
|
||||
mock_notification_handler_class.return_value = mock_notification_handler
|
||||
|
||||
mock_activities = AsyncMock()
|
||||
mock_activities.shutdown = AsyncMock()
|
||||
mock_activities_class.return_value = mock_activities
|
||||
|
||||
mock_app_up = Mock()
|
||||
mock_metrics.APP_UP.labels.return_value = mock_app_up
|
||||
|
||||
mock_plugin_store_instance = AsyncMock()
|
||||
mock_plugin_store_instance.install_runtime = AsyncMock(
|
||||
return_value={'runtime': 'model-manager-worker', 'installed': []},
|
||||
)
|
||||
mock_plugin_store_class.return_value = mock_plugin_store_instance
|
||||
mock_build_plugin_store_config.return_value = {
|
||||
'base_url': 'http://sientia-plugin-store.svc.cluster.local',
|
||||
'owner': 'sientia',
|
||||
'repo': 'model-library-store',
|
||||
'branch': 'main',
|
||||
'username': 'gitea-user',
|
||||
'password': 'gitea-password',
|
||||
'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000',
|
||||
'pypi_username': None,
|
||||
'pypi_password': None,
|
||||
}
|
||||
|
||||
# Ensure RUNTIME is not defined
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
await main()
|
||||
|
||||
assert exc_info.value.code == 1
|
||||
mock_logger.custom_critical.assert_called_once()
|
||||
mock_app_up.set.assert_called_with(0)
|
||||
|
||||
# Run main() - should not fail despite schedule creation error
|
||||
with pytest.raises(SystemExit):
|
||||
await main()
|
||||
|
||||
392
values.yaml
392
values.yaml
@@ -1,26 +1,222 @@
|
||||
# Default values for sientia-module.
|
||||
#
|
||||
# Default values for sientia-dataops-model-manager using the sientia-module chart (0.6.x).
|
||||
# This is a YAML-formatted file.
|
||||
# Declare variables to be passed into your templates.
|
||||
#
|
||||
|
||||
# This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/
|
||||
replicaCount: 1
|
||||
projectName: &projectName "sientia-dataops-model-manager"
|
||||
|
||||
# This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/
|
||||
# -----------------------------------------------------------------------------
|
||||
# Image configuration (chart-level)
|
||||
# -----------------------------------------------------------------------------
|
||||
# The sientia-module chart allows overriding the image used by all runtimes.
|
||||
# Per requirement, we deploy using the sientia-module image v1.0.0.
|
||||
image:
|
||||
repository: aignosi.azurecr.io/sientia-dataops-model-manager
|
||||
# This sets the pull policy for images.
|
||||
repository: aignosi.azurecr.io/sientia-module
|
||||
pullPolicy: IfNotPresent
|
||||
# Overrides the image tag whose default is the chart appVersion.
|
||||
tag: "1.1.2"
|
||||
tag: "1.0.0"
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Global configuration shared by all runtimes
|
||||
# -----------------------------------------------------------------------------
|
||||
global:
|
||||
# Namespace used by the chart.
|
||||
namespace: sientia
|
||||
|
||||
# Common labels applied to pods (can be extended per project).
|
||||
commonLabels: {}
|
||||
|
||||
# Resources inherited by all runtimes unless overridden.
|
||||
resources: {}
|
||||
# We usually recommend not to specify default resources and to leave this as a conscious
|
||||
# choice for the user. This also increases chances charts run on environments with little
|
||||
# resources, such as Minikube. If you do want to specify resources, uncomment the following
|
||||
# lines, adjust them as necessary, and remove the curly braces after 'resources:'.
|
||||
# limits:
|
||||
# cpu: 100m
|
||||
# memory: 128Mi
|
||||
# requests:
|
||||
# cpu: 100m
|
||||
# memory: 128Mi
|
||||
|
||||
# Probes inherited by all runtimes unless overridden.
|
||||
# More information:
|
||||
# https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/
|
||||
livenessProbe:
|
||||
exec:
|
||||
command:
|
||||
- python3
|
||||
- -c
|
||||
- "import requests; requests.get('http://localhost:9090/metrics')"
|
||||
initialDelaySeconds: 20
|
||||
periodSeconds: 30
|
||||
|
||||
readinessProbe:
|
||||
exec:
|
||||
command:
|
||||
- python3
|
||||
- -c
|
||||
- "import requests; requests.get('http://localhost:9090/metrics')"
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
|
||||
# Autoscaling configuration inherited by all runtimes unless overridden.
|
||||
# More information:
|
||||
# https://kubernetes.io/docs/concepts/workloads/autoscaling/
|
||||
autoscaling:
|
||||
enabled: false
|
||||
minReplicas: 1
|
||||
maxReplicas: 100
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# targetMemoryUtilizationPercentage: 80
|
||||
|
||||
# Environment variables shared by all runtimes.
|
||||
env:
|
||||
- name: POSTGRES_HOST
|
||||
value: "paradedb-rw.paradedb.svc.cluster.local"
|
||||
- name: POSTGRES_PORT
|
||||
value: "5432"
|
||||
- name: POSTGRES_USER
|
||||
value: "postgres"
|
||||
- name: POSTGRES_PASSWORD
|
||||
value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3"
|
||||
- name: POSTGRES_DBNAME
|
||||
value: "sientia-core-mlops-bff"
|
||||
- name: POSTGRES_MIN_CONNECTIONS
|
||||
value: "10"
|
||||
- name: POSTGRES_MAX_CONNECTIONS
|
||||
value: "30"
|
||||
|
||||
- name: MLFLOW_URL
|
||||
value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local:80"
|
||||
- name: MLFLOW_USERNAME
|
||||
value: "aignosi"
|
||||
- name: MLFLOW_PASSWORD
|
||||
value: "1L0FP50j3ncp123"
|
||||
|
||||
- name: LOG_LEVEL
|
||||
value: "DEBUG"
|
||||
- name: HTTP_METRICS_PORT
|
||||
value: "9090"
|
||||
- name: HTTP_SDK_METRICS_PORT
|
||||
value: "9091"
|
||||
|
||||
- name: TEMPORAL_HOST
|
||||
value: "temporal-frontend.temporal.svc.cluster.local:7233"
|
||||
- name: TEMPORAL_NAMESPACE
|
||||
value: "model-manager"
|
||||
- name: TRAIN_TASK_QUEUE
|
||||
value: "train_model-queue"
|
||||
- name: CLEANUP_TASK_QUEUE
|
||||
value: "cleanup-queue"
|
||||
- name: TEMPORAL_USE_TLS
|
||||
value: "false"
|
||||
|
||||
- name: STORE_BASE_URL
|
||||
value: "http://gitea-http.gitea.svc.cluster.local"
|
||||
- name: STORE_OWNER
|
||||
value: "aignosi"
|
||||
- name: STORE_REPO
|
||||
value: "suse-model-store"
|
||||
- name: STORE_USERNAME
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: sientia-plugin-store-credentials
|
||||
key: username
|
||||
- name: STORE_PASSWORD
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: sientia-plugin-store-credentials
|
||||
key: password
|
||||
- name: STORE_CACHE_TTL_SECONDS
|
||||
value: "3600"
|
||||
|
||||
- name: MONGODB_USERNAME
|
||||
value: "root"
|
||||
- name: MONGODB_PASSWORD
|
||||
value: "wKZDbMNU1c"
|
||||
- name: MONGODB_URL
|
||||
value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
|
||||
- name: MONGODB_DATABASE
|
||||
value: "sientia"
|
||||
- name: MONGODB_TTL_INDEX_HOURS
|
||||
value: "1"
|
||||
|
||||
- name: MINIO_ENDPOINT_URL
|
||||
value: "http://minio.minio.svc.cluster.local:9000"
|
||||
- name: MINIO_ACCESS_KEY
|
||||
value: "model-training-user"
|
||||
- name: MINIO_SECRET_KEY
|
||||
value: "modelTrainingUser123"
|
||||
- name: MINIO_REGION
|
||||
value: "us-east-1"
|
||||
- name: MINIO_USE_SSL
|
||||
value: "false"
|
||||
- name: MINIO_MAX_RETRY_ATTEMPTS
|
||||
value: "3"
|
||||
- name: MINIO_RETRY_MODE
|
||||
value: "adaptive"
|
||||
- name: MINIO_CONNECT_TIMEOUT
|
||||
value: "10"
|
||||
- name: MINIO_READ_TIMEOUT
|
||||
value: "60"
|
||||
|
||||
- name: TIMEOUT_VALIDATE_PARAMS
|
||||
value: "30"
|
||||
- name: TIMEOUT_TRAIN_MODEL
|
||||
value: "2700"
|
||||
- name: TIMEOUT_DELETE_FILE
|
||||
value: "120"
|
||||
- name: TIMEOUT_UPDATE_DATABASE
|
||||
value: "30"
|
||||
|
||||
- name: CLEANUP_RETENTION_HOURS
|
||||
value: "24"
|
||||
- name: CLEANUP_DRY_RUN
|
||||
value: "false"
|
||||
- name: TIMEOUT_CLEANUP_MINIO
|
||||
value: "300"
|
||||
- name: TIMEOUT_CLEANUP_LOCAL
|
||||
value: "120"
|
||||
- name: MAX_KEYS_CLEANUP
|
||||
value: "1000"
|
||||
- name: DEFAULT_CLEANUP_BUCKET
|
||||
value: "model-training"
|
||||
|
||||
# Cleanup Schedule Configuration
|
||||
- name: CLEANUP_SCHEDULE_ID
|
||||
value: "cleanup-files-daily"
|
||||
- name: CLEANUP_CRON
|
||||
value: "0 0 * * *" # Midnight UTC
|
||||
- name: CLEANUP_TIMEZONE
|
||||
value: "UTC"
|
||||
- name: CLEANUP_EXECUTION_TIMEOUT_HOURS
|
||||
value: "1"
|
||||
|
||||
- name: PYPI_SERVER
|
||||
value: "http://library-distribution-server.library.svc.cluster.local:5000"
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Runtimes configuration
|
||||
# -----------------------------------------------------------------------------
|
||||
# Each runtime inherits settings from `global` (resources, env, probes, autoscaling)
|
||||
# unless overridden here.
|
||||
runtimes:
|
||||
- name: "model-manager-worker"
|
||||
# Replicas for this runtime. Replaces the old replicaCount.
|
||||
replicas: 1
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Chart-level configuration (applies to all runtimes)
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
# This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/
|
||||
imagePullSecrets:
|
||||
- name: docker-hub-secret
|
||||
- name: docker-hub-secret
|
||||
|
||||
# This is to override the chart name.
|
||||
nameOverride: "sientia-dataops-model-manager"
|
||||
fullnameOverride: "sientia-dataops-model-manager"
|
||||
namespace: sientia
|
||||
nameOverride: *projectName
|
||||
fullnameOverride: *projectName
|
||||
|
||||
# This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/
|
||||
serviceAccount:
|
||||
@@ -32,7 +228,7 @@ serviceAccount:
|
||||
annotations: {}
|
||||
# The name of the service account to use.
|
||||
# If not set and create is true, a name is generated using the fullname template
|
||||
name: "sientia-dataops-model-manager"
|
||||
name: *projectName
|
||||
|
||||
# This is for setting Kubernetes Annotations to a Pod.
|
||||
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
|
||||
@@ -52,45 +248,6 @@ securityContext: {}
|
||||
# runAsNonRoot: true
|
||||
# runAsUser: 1000
|
||||
|
||||
resources: {}
|
||||
# We usually recommend not to specify default resources and to leave this as a conscious
|
||||
# choice for the user. This also increases chances charts run on environments with little
|
||||
# resources, such as Minikube. If you do want to specify resources, uncomment the following
|
||||
# lines, adjust them as necessary, and remove the curly braces after 'resources:'.
|
||||
# limits:
|
||||
# cpu: 100m
|
||||
# memory: 128Mi
|
||||
# requests:
|
||||
# cpu: 100m
|
||||
# memory: 128Mi
|
||||
|
||||
# This is to setup the liveness and readiness probes more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/
|
||||
livenessProbe:
|
||||
exec:
|
||||
command:
|
||||
- python3
|
||||
- -c
|
||||
- "import requests; requests.get('http://localhost:9090/metrics')"
|
||||
initialDelaySeconds: 20
|
||||
periodSeconds: 30
|
||||
|
||||
readinessProbe:
|
||||
exec:
|
||||
command:
|
||||
- python3
|
||||
- -c
|
||||
- "import requests; requests.get('http://localhost:9090/metrics')"
|
||||
initialDelaySeconds: 10
|
||||
periodSeconds: 15
|
||||
|
||||
# This section is for setting up autoscaling more information can be found here: https://kubernetes.io/docs/concepts/workloads/autoscaling/
|
||||
autoscaling:
|
||||
enabled: false
|
||||
minReplicas: 1
|
||||
maxReplicas: 100
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# targetMemoryUtilizationPercentage: 80
|
||||
|
||||
# Additional volumes on the output Deployment definition.
|
||||
volumes:
|
||||
- name: reports-volume
|
||||
@@ -151,119 +308,6 @@ serviceMonitor:
|
||||
additionalLabels:
|
||||
release: kube-prometheus-stack
|
||||
|
||||
env:
|
||||
- name: POSTGRES_HOST
|
||||
value: "paradedb-rw.paradedb.svc.cluster.local"
|
||||
- name: POSTGRES_PORT
|
||||
value: "5432"
|
||||
- name: POSTGRES_USER
|
||||
value: "postgres"
|
||||
- name: POSTGRES_PASSWORD
|
||||
value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3"
|
||||
- name: POSTGRES_DBNAME
|
||||
value: "sientia-core-mlops-bff"
|
||||
- name: POSTGRES_MIN_CONNECTIONS
|
||||
value: "10"
|
||||
- name: POSTGRES_MAX_CONNECTIONS
|
||||
value: "30"
|
||||
|
||||
- name: MLFLOW_URL
|
||||
value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local:80"
|
||||
- name: MLFLOW_USERNAME
|
||||
value: "aignosi"
|
||||
- name: MLFLOW_PASSWORD
|
||||
value: "1L0FP50j3ncp123"
|
||||
|
||||
- name: LOG_LEVEL
|
||||
value: "DEBUG"
|
||||
- name: HTTP_METRICS_PORT
|
||||
value: "9090"
|
||||
- name: HTTP_SDK_METRICS_PORT
|
||||
value: "9091"
|
||||
- name: PROJECT_NAME
|
||||
value: "sientia-model-manager"
|
||||
|
||||
- name: TEMPORAL_HOST
|
||||
value: "temporal-frontend.temporal.svc.cluster.local:7233"
|
||||
- name: TEMPORAL_NAMESPACE
|
||||
value: "model-manager"
|
||||
- name: TRAIN_TASK_QUEUE
|
||||
value: "train_model-queue"
|
||||
- name: CLEANUP_TASK_QUEUE
|
||||
value: "cleanup-queue"
|
||||
- name: TEMPORAL_USE_TLS
|
||||
value: "false"
|
||||
|
||||
- name: MONGODB_USERNAME
|
||||
value: "root"
|
||||
- name: MONGODB_PASSWORD
|
||||
value: "wKZDbMNU1c"
|
||||
- name: MONGODB_URL
|
||||
value: "my-release-mongodb.mongodb.svc.cluster.local:27017"
|
||||
- name: MONGODB_DATABASE
|
||||
value: "sientia"
|
||||
- name: MONGODB_TTL_INDEX_HOURS
|
||||
value: "1"
|
||||
|
||||
- name: MINIO_ENDPOINT_URL
|
||||
value: "http://minio.minio.svc.cluster.local:9000"
|
||||
- name: MINIO_ACCESS_KEY
|
||||
value: "model-training-user"
|
||||
- name: MINIO_SECRET_KEY
|
||||
value: "modelTrainingUser123"
|
||||
- name: MINIO_REGION
|
||||
value: "us-east-1"
|
||||
- name: MINIO_USE_SSL
|
||||
value: "false"
|
||||
- name: MINIO_MAX_RETRY_ATTEMPTS
|
||||
value: "3"
|
||||
- name: MINIO_RETRY_MODE
|
||||
value: "adaptive"
|
||||
- name: MINIO_CONNECT_TIMEOUT
|
||||
value: "10"
|
||||
- name: MINIO_READ_TIMEOUT
|
||||
value: "60"
|
||||
|
||||
- name: TIMEOUT_VALIDATE_PARAMS
|
||||
value: "30"
|
||||
- name: TIMEOUT_TRAIN_MODEL
|
||||
value: "2700"
|
||||
- name: TIMEOUT_DELETE_FILE
|
||||
value: "120"
|
||||
- name: TIMEOUT_UPDATE_DATABASE
|
||||
value: "30"
|
||||
|
||||
- name: CLEANUP_RETENTION_HOURS
|
||||
value: "24"
|
||||
- name: CLEANUP_DRY_RUN
|
||||
value: "false"
|
||||
- name: TIMEOUT_CLEANUP_MINIO
|
||||
value: "300"
|
||||
- name: TIMEOUT_CLEANUP_LOCAL
|
||||
value: "120"
|
||||
- name: MAX_KEYS_CLEANUP
|
||||
value: "1000"
|
||||
- name: DEFAULT_CLEANUP_BUCKET
|
||||
value: "model-training"
|
||||
|
||||
# Cleanup Schedule Configuration
|
||||
- name: CLEANUP_SCHEDULE_ID
|
||||
value: "cleanup-files-daily"
|
||||
- name: CLEANUP_CRON
|
||||
value: "0 0 * * *" # Midnight UTC
|
||||
- name: CLEANUP_TIMEZONE
|
||||
value: "UTC"
|
||||
- name: CLEANUP_EXECUTION_TIMEOUT_HOURS
|
||||
value: "1"
|
||||
|
||||
- name: EXTRA_PIP_REQUIREMENTS
|
||||
value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git"
|
||||
|
||||
- name: POD_ID
|
||||
valueFrom:
|
||||
fieldRef:
|
||||
fieldPath: metadata.name
|
||||
|
||||
ssh:
|
||||
enabled: false
|
||||
secretName: git-ssh-key-sientia-model-manager-worker
|
||||
@@ -308,11 +352,21 @@ grafanaDatasource:
|
||||
# jsonData:
|
||||
# timeInterval: "5s"
|
||||
|
||||
# kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp
|
||||
|
||||
# helm upgrade --install sientia-dataops-model-manager sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.0
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Helm usage examples
|
||||
# -----------------------------------------------------------------------------
|
||||
# kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=<pwd>
|
||||
#
|
||||
# helm upgrade --install sientia-dataops-model-manager sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.1
|
||||
#
|
||||
# Global/runtimes layout note:
|
||||
# - Shared configuration lives under `global` (env, probes, autoscaling, namespace).
|
||||
# - Individual runtimes are defined under `runtimes`, each with its own `name` and `replicas`.
|
||||
# - Runtimes inherit `global` settings unless overridden at the runtime level.
|
||||
#
|
||||
# kubectl create secret generic git-ssh-key-sientia-model-manager-worker \
|
||||
# --namespace sientia \
|
||||
# --from-file=ssh-privatekey=git_key \
|
||||
# --type=kubernetes.io/ssh-auth
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user