gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: integrate PluginStore and MinIO repository into model manager activities

Browse Source 
- 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
2026-03-11 17:35:05 -03:00
parent 9d71c0cf80
commit cf5111e520
23 changed files with 1480 additions and 4588 deletions
Download Patch File Download Diff File Expand all files Collapse all files

385
model-manager-plugin-store-migration-plan.md Normal file
View 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]]

52
model_manager/activities/activities.py
View File

@@ -6,12 +6,14 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController 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.cleanup import Cleanup
from model_manager.activities.experiment_tracking import ExperimentTracking from model_manager.activities.experiment_tracking import ExperimentTracking
from model_manager.activities.training import Training 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): class Activities(ExperimentTracking, Training, Cleanup):
@@ -40,8 +42,10 @@ class Activities(ExperimentTracking, Training, Cleanup):
postgres_config: dict[str, Any], postgres_config: dict[str, Any],
mlflow_config: dict[str, Any], mlflow_config: dict[str, Any],
minio_config: dict[str, Any], minio_config: dict[str, Any],
plugin_store: PluginStore,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController,
): ):
""" """
Initialize the Activities orchestrator with all required configurations. Initialize the Activities orchestrator with all required configurations.
@@ -62,9 +66,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
Raises: Raises:
Exception: If any parent class initialization fails Exception: If any parent class initialization fails
""" """
metrics_controller = MetricsController(
logger=logger,
)
ExperimentTracking.__init__( ExperimentTracking.__init__(
self, self,
@@ -80,30 +82,40 @@ class Activities(ExperimentTracking, Training, Cleanup):
metrics_controller=metrics_controller, metrics_controller=metrics_controller,
) )
self.model_repository = ModelRepository( self.mlflow_repository = SientiaMLflowRepository(
url=mlflow_config['url'], host=mlflow_config['url'],
username=mlflow_config['username'], username=mlflow_config['username'],
password=mlflow_config['password'], password=mlflow_config['password'],
logger=logger, logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
self.storage_repository = StorageRepository( # MinIO repository used for all object storage operations
endpoint_url=minio_config['endpoint_url'], 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'], access_key=minio_config['access_key'],
secret_key=minio_config['secret_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, logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,
secure=minio_config['use_ssl'],
) )
Training.__init__( Training.__init__(
self, self,
model_repository=self.model_repository, mlflow_repository=self.mlflow_repository,
storage_repository=self.storage_repository, plugin_store=plugin_store,
minio_repository=self.minio_repository,
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller, metrics_controller=metrics_controller,
@@ -111,7 +123,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
Cleanup.__init__( Cleanup.__init__(
self, self,
storage_repository=self.storage_repository, minio_repository=self.minio_repository,
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller, metrics_controller=metrics_controller,
@@ -137,7 +149,7 @@ class Activities(ExperimentTracking, Training, Cleanup):
# Logging here could cause issues if logger is already destroyed # Logging here could cause issues if logger is already destroyed
pass pass
async def shutdown(self): def shutdown(self):
""" """
Gracefully shutdown all activities and clean up resources. Gracefully shutdown all activities and clean up resources.
@@ -152,4 +164,4 @@ class Activities(ExperimentTracking, Training, Cleanup):
""" """
ExperimentTracking.close(self) ExperimentTracking.close(self)
self.info('Postgres client closed') self.info('Postgres client closed')
self.storage_repository.close() SientiaMonitoring.shutdown(self)

32
model_manager/activities/cleanup.py
View File

@@ -21,9 +21,9 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.logger import Logger from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring 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.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')) RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24'))
DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true' DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true'
@@ -41,7 +41,7 @@ class Cleanup(SientiaMonitoring):
def __init__( def __init__(
self, self,
storage_repository: StorageRepository, minio_repository: MinioRepository,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController, metrics_controller: MetricsController,
@@ -50,13 +50,13 @@ class Cleanup(SientiaMonitoring):
Initialize Cleanup activity. Initialize Cleanup activity.
Args: Args:
storage_repository: Repository for MinIO operations minio_repository: Repository for MinIO operations
logger: Logger instance for observability logger: Logger instance for observability
notification_handler: Handler for sending notifications notification_handler: Handler for sending notifications
metrics_controller: Controller for metrics emission metrics_controller: Controller for metrics emission
""" """
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.storage_repository = storage_repository self.minio_repository = minio_repository
# Configuration from environment variables # Configuration from environment variables
self.retention_hours = RETENTION_HOURS self.retention_hours = RETENTION_HOURS
@@ -113,15 +113,21 @@ class Cleanup(SientiaMonitoring):
files_deleted = 0 files_deleted = 0
errors = [] errors = []
# List objects in the specified bucket # List objects in the specified bucket. MinioRepository applies BASE_PREFIX
max_keys = self.max_keys_cleanup # Use environment variable for page size # internally; we request all objects under that prefix for this bucket.
objects = self.storage_repository.list_bucket_objects(bucket_name, max_keys) objects = await self.minio_repository.list_objects(
prefix='',
bucket=bucket_name,
recursive=True,
metadata=metadata,
)
for obj_key in objects: for obj_key in objects:
files_scanned += 1 files_scanned += 1
# Extract timestamp from filename # Extract timestamp from the filename portion of the object key
match = self.minio_timestamp_pattern.match(obj_key) filename = obj_key.split('/')[-1]
match = self.minio_timestamp_pattern.match(filename)
if not match: if not match:
self.debug(f'Skipping file without timestamp pattern: {obj_key}', metadata) self.debug(f'Skipping file without timestamp pattern: {obj_key}', metadata)
continue continue
@@ -137,10 +143,14 @@ class Cleanup(SientiaMonitoring):
files_deleted += 1 files_deleted += 1
else: else:
try: 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) self.info(f'Deleted stale file: {obj_key}', metadata)
files_deleted += 1 files_deleted += 1
except OSError as e: except Exception as e: # noqa: BLE001
error_msg = f'Failed to delete {obj_key}: {str(e)}' error_msg = f'Failed to delete {obj_key}: {str(e)}'
errors.append(error_msg) errors.append(error_msg)
self.error(error_msg, metadata) self.error(error_msg, metadata)

116
model_manager/activities/training.py
View File

@@ -12,18 +12,20 @@ with workflow.unsafe.imports_passed_through():
import traceback import traceback
from typing import Any from typing import Any
import pandas as pd
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.notifications.models import NotificationLevel from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.logger import Logger from sientia_do.observability.logger import Logger
from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.observability.sientia_monitoring import SientiaMonitoring 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.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
from model_manager.utils.exceptions import ModelTrainingError from model_manager.utils.exceptions import ModelTrainingError
from model_manager.utils.models.train_model_params import TrainModelParams from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.repository.model_repository import ModelRepository from model_manager.utils.repository.data_manager_repository import DataManagerRepository
from model_manager.utils.repository.storage_repository import StorageRepository
from model_manager.utils.repository.training_repository import TrainingRepository
class Training(SientiaMonitoring): class Training(SientiaMonitoring):
@@ -38,8 +40,9 @@ class Training(SientiaMonitoring):
def __init__( def __init__(
self, self,
model_repository: ModelRepository, mlflow_repository: SientiaMLflowRepository,
storage_repository: StorageRepository, plugin_store: PluginStore,
minio_repository: MinioRepository,
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, notification_handler: NotificationHandler,
metrics_controller: MetricsController, metrics_controller: MetricsController,
@@ -52,9 +55,10 @@ class Training(SientiaMonitoring):
notification_handler: Handler for sending notifications notification_handler: Handler for sending notifications
""" """
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.training_repository = TrainingRepository(logger) self.data_manager_repository = DataManagerRepository(logger)
self.model_repository = model_repository self.mlflow_repository = mlflow_repository
self.storage_repository = storage_repository self.plugin_store = plugin_store
self.minio_repository = minio_repository
@activity.defn(name='validate_train_params') @activity.defn(name='validate_train_params')
async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams: 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: This activity orchestrates the ML training pipeline:
1. Validate input parameters. 1. Validate input parameters.
2. Train the model via TrainingRepository. 2. Prepare data via DataManagerRepository.
3. Perform post-training calculations. 3. Train the model and compute metrics.
Args: Args:
input_data: Training configuration containing: input_data: Training configuration containing:
@@ -130,40 +134,97 @@ class Training(SientiaMonitoring):
- train_params (TrainModelParams | dict): Training parameters. - train_params (TrainModelParams | dict): Training parameters.
Returns: Returns:
dict: Keys `run_name` and `run_dir` when training and saving succeed. dict: Key `run_name` when training and saving succeed.
Raises: Raises:
ValueError: If input validation fails. ValueError: If input validation fails.
Exception: If training fails (after sending notification). Exception: If training fails (after sending notification).
""" """
metadata = input_data.get('metadata', {}) metadata = input_data.get('metadata')
train_params = input_data['train_params'] train_params = input_data['train_params']
if isinstance(train_params, dict): if isinstance(train_params, dict):
train_params = TrainModelParams.from_dict(train_params) train_params = TrainModelParams.from_dict(train_params)
# type: ignore[assignment]
model_trained = False model_trained = False
model_saved = False model_saved = False
metrics_status = 'success' metrics_status = 'success'
try: try:
with self.storage_repository.fetch_file( # Download training file bytes from MinIO
train_params.bucket_name, train_params.file_name train_bytes = await self.minio_repository.download_file(
) as uploaded_file: object_name=train_params.file_name,
train_result = self.training_repository.train(uploaded_file, train_params) bucket=train_params.bucket_name,
metadata=metadata,
)
train_result = self.training_repository.after_train_calculation( # Download optional validation file bytes from the same bucket
train_params, train_result 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,
)
train_result = self.data_manager_repository.prepare_training_data(
train_file_bytes=train_bytes,
validation_file_bytes=val_bytes,
params=train_params,
metadata=metadata,
)
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 model_trained = True
train_result = self.model_repository.save_model(train_result)
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 model_saved = True
return { return {
'run_name': train_result.run_name, 'run_name': run_info.run_name or run_info.run_id,
'run_dir': train_result.run_dir,
} }
except Exception as e: # noqa: BLE001 except Exception as e: # noqa: BLE001
metrics_status = 'error' metrics_status = 'error'
@@ -205,7 +266,6 @@ class Training(SientiaMonitoring):
Args: Args:
input_data: Cleanup configuration containing: input_data: Cleanup configuration containing:
- metadata (dict): Workflow execution metadata. - metadata (dict): Workflow execution metadata.
- run_dir (str): Temporary directory to remove.
- bucket_name (str): MinIO bucket of the uploaded file. - bucket_name (str): MinIO bucket of the uploaded file.
- file_name (str): MinIO object key to delete. - file_name (str): MinIO object key to delete.
@@ -213,19 +273,21 @@ class Training(SientiaMonitoring):
Exception: If cleanup fails (after sending notification). Exception: If cleanup fails (after sending notification).
""" """
metadata = input_data.get('metadata', {}) metadata = input_data.get('metadata', {})
run_dir = input_data.get('run_dir', '')
bucket_name = input_data.get('bucket_name', '') bucket_name = input_data.get('bucket_name', '')
file_name = input_data.get('file_name', '') file_name = input_data.get('file_name', '')
metrics_status = 'success' metrics_status = 'success'
try: try:
self.model_repository.cleanup_run_directory(run_dir) await self.minio_repository.delete_file(
self.storage_repository.delete_file(bucket_name, file_name) object_name=file_name,
bucket=bucket_name,
metadata=metadata,
)
except Exception as e: # noqa: BLE001 except Exception as e: # noqa: BLE001
metrics_status = 'error' metrics_status = 'error'
error_msg = ( error_msg = (
f'Error cleaning up resources - Run directory: {run_dir}, ' 'Error cleaning up resources - '
f'File: {bucket_name}/{file_name}, Error: {str(e)}' f'File: {bucket_name}/{file_name}, Error: {str(e)}'
) )

40
model_manager/utils/connectors_config.py
View File

@@ -121,3 +121,43 @@ def build_minio_config() -> dict[str, Any]:
'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')), 'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')),
'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')), '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'),
}

26
model_manager/utils/models/train_model_params.py
View File

@@ -33,7 +33,8 @@ class TrainModelParams:
use_scaler (bool): Whether to use a scaler for data normalization. use_scaler (bool): Whether to use a scaler for data normalization.
include_ar (bool): Whether to include autoregressive variables. include_ar (bool): Whether to include autoregressive variables.
bucket_name (str): Name of the MinIO bucket containing training data. 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. line_separator (str): Line separator used in the CSV file.
decimal_separator (str): Decimal 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. 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 include_ar: bool
bucket_name: str bucket_name: str
file_name: str file_name: str
validation_file_name: str | None
line_separator: str line_separator: str
decimal_separator: str decimal_separator: str
date_column: str | None date_column: str | None
@@ -122,6 +124,9 @@ class TrainModelParams:
include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'), include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'),
bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'), bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'),
file_name=cls._check_none(data.get('file_name'), str, 'file_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'), line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'),
decimal_separator=cls._check_none( decimal_separator=cls._check_none(
data.get('decimal_separator'), str, 'decimal_separator' data.get('decimal_separator'), str, 'decimal_separator'
@@ -220,6 +225,7 @@ class TrainModelParams:
self._validate_limits() self._validate_limits()
self._validate_required_strings() self._validate_required_strings()
self._validate_date_format() self._validate_date_format()
self._validate_validation_file_name()
def _validate_numeric_ranges(self) -> None: def _validate_numeric_ranges(self) -> None:
"""Validate numeric parameters are within acceptable ranges.""" """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.""" """Validate date_format is one of the allowed frontend formats when set."""
if self.date_format: if self.date_format:
validate_frontend_date_format(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')

11
model_manager/utils/models/train_model_result.py
View File

@@ -2,7 +2,6 @@ from dataclasses import dataclass
import pandas as pd import pandas as pd
from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel
from model_manager.utils.models.train_model_params import TrainModelParams 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. A data container for storing the results of a machine learning training process.
This dataclass encapsulates all outputs from the training pipeline, including 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. It is used to pass results between activities in the training workflow.
Attributes: Attributes:
params (TrainModelParams): The parameters used to train the model. 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_train (pd.DataFrame): The training dataset features.
x_test (pd.DataFrame): The testing dataset features. x_test (pd.DataFrame): The testing dataset features.
y_train (pd.DataFrame): The training dataset target values. y_train (pd.DataFrame): The training dataset target values.
y_test (pd.DataFrame): The testing 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_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. 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. 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. 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. 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. 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 params: TrainModelParams
process_data: DataPreprocessor
x_train: pd.DataFrame x_train: pd.DataFrame
x_test: pd.DataFrame x_test: pd.DataFrame
y_train: pd.Series y_train: pd.Series
y_test: pd.Series y_test: pd.Series
regr: LinearRegressionModel
scaler_dict: dict
y_pred: pd.Series | None = None y_pred: pd.Series | None = None
y_train_pred: pd.Series | None = None y_train_pred: pd.Series | None = None
mse_val: float | None = None mse_val: float | None = None
@@ -57,4 +49,3 @@ class TrainModelResult:
report_path: str | None = None report_path: str | None = None
train_data_path: str | None = None train_data_path: str | None = None
test_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
View 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

471
model_manager/utils/repository/model_repository.py
View File

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

164
model_manager/utils/repository/storage_repository.py
View File

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

519
model_manager/utils/repository/training_repository.py
View File

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

79
model_manager/worker/worker.py
View File

@@ -40,6 +40,9 @@ with workflow.unsafe.imports_passed_through():
from prometheus_client import start_http_server from prometheus_client import start_http_server
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
from sientia_do.observability.logger import Logger as SientiaLogger 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 import metrics
from model_manager.activities.activities import Activities from model_manager.activities.activities import Activities
@@ -48,6 +51,7 @@ with workflow.unsafe.imports_passed_through():
build_minio_config, build_minio_config,
build_mlflow_config, build_mlflow_config,
build_mongodb_config, build_mongodb_config,
build_plugin_store_config,
build_postgres_config, build_postgres_config,
) )
from model_manager.utils.logger_helper import get_logger 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 from model_manager.workflows.train_model import TrainModel
POD_ID = os.getenv('POD_ID') POD_ID = os.getenv('POD_ID')
RUNTIME = os.getenv('RUNTIME')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))
TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue') TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue')
CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue') CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue')
@@ -79,14 +84,20 @@ async def main():
Exception: Any unhandled exception during worker execution Exception: Any unhandled exception during worker execution
SystemExit: On graceful shutdown or error conditions SystemExit: On graceful shutdown or error conditions
""" """
if not RUNTIME:
raise ValueError('RUNTIME environment variable is required')
host = os.getenv('TEMPORAL_HOST', 'localhost:7233') host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true' use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
logger = get_logger(__name__) logger = get_logger(__name__)
metadata = { metadata = {
'pod_id': POD_ID, 'pod_id': POD_ID,
'runtime': RUNTIME,
} }
start_prometheus_server(logger, metadata) start_prometheus_server(logger, metadata)
mongo_config = build_mongodb_config() 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'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( activities = Activities(
postgres_config=build_postgres_config(), postgres_config=build_postgres_config(),
mlflow_config=build_mlflow_config(), mlflow_config=build_mlflow_config(),
minio_config=build_minio_config(), minio_config=build_minio_config(),
plugin_store=plugin_store,
logger=logger, logger=logger,
notification_handler=notification_handler, notification_handler=notification_handler,
metrics_controller=metrics_controller,
) )
new_runtime = Runtime( new_runtime = Runtime(
@@ -133,44 +173,33 @@ async def main():
# The schedule can be created manually if needed # The schedule can be created manually if needed
workers = [ workers = [
Worker( prepare_worker(
temporal_client, main_workflow=TrainModel,
task_queue=TRAIN_TASK_QUEUE, other_workflows=[],
workflows=[TrainModel],
activities=[ activities=[
activities.update_experiment_run, activities.update_experiment_run,
activities.validate_train_params, activities.validate_train_params,
activities.train_model, activities.train_model,
activities.cleanup_resources, activities.cleanup_resources,
], ],
max_concurrent_workflow_tasks=10, temporal_client=temporal_client,
max_concurrent_activities=10, logger=logger,
max_concurrent_local_activities=10,
max_cached_workflows=100,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
), ),
Worker( prepare_worker(
temporal_client, main_workflow=CleanupFiles,
task_queue=CLEANUP_TASK_QUEUE, other_workflows=[],
workflows=[CleanupFiles],
activities=[ activities=[
activities.cleanup_minio_files, activities.cleanup_minio_files,
activities.cleanup_temp_directories, activities.cleanup_temp_directories,
], ],
max_concurrent_workflow_tasks=20, temporal_client=temporal_client,
max_concurrent_activities=20, logger=logger,
max_concurrent_local_activities=20,
max_cached_workflows=100,
workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
activity_task_poller_behavior=PollerBehaviorAutoscaling(),
), ),
] ]
handlers = [] handlers = [
w.run() for w in workers
for w in workers: ]
handlers.append(w.run())
logger.custom_info('Model manager workers initialized', metadata) logger.custom_info('Model manager workers initialized', metadata)
@@ -183,7 +212,7 @@ async def main():
finally: finally:
notification_handler.shutdown() notification_handler.shutdown()
logger.custom_info('MongoDB client closed', metadata) logger.custom_info('MongoDB client closed', metadata)
await activities.shutdown() activities.shutdown()
# Exit with a non-zero status code to indicate failure to Kubernetes # 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 metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN
sys.exit(1) sys.exit(1)

8
model_manager/workflows/train_model.py
View File

@@ -118,7 +118,6 @@ class TrainModel:
await self._cleanup_resources( await self._cleanup_resources(
experiment_run_id=experiment_run_id, experiment_run_id=experiment_run_id,
run_dir=(train_result.get('run_dir') or ''),
bucket_name=train_params.bucket_name, bucket_name=train_params.bucket_name,
file_name=train_params.file_name, file_name=train_params.file_name,
metadata=metadata, metadata=metadata,
@@ -275,7 +274,6 @@ class TrainModel:
async def _cleanup_resources( async def _cleanup_resources(
self, self,
experiment_run_id: int, experiment_run_id: int,
run_dir: str,
bucket_name: str, bucket_name: str,
file_name: str, file_name: str,
metadata: dict[str, Any], metadata: dict[str, Any],
@@ -283,12 +281,11 @@ class TrainModel:
""" """
Cleanup resources and delete file from MinIO. Cleanup resources and delete file from MinIO.
This method removes the temporary run directory via activity and deletes This method deletes the training file from MinIO. On success, updates
the training file from MinIO. On success, updates DB status to FILE_DELETED. DB status to FILE_DELETED.
On error, updates DB status to FILE_DELETE_ERROR. On error, updates DB status to FILE_DELETE_ERROR.
Args: Args:
saved_result: TrainModelResult with run_dir and params information
experiment_run_id: Validated experiment run ID experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata metadata: Workflow execution metadata
@@ -300,7 +297,6 @@ class TrainModel:
Activities.cleanup_resources, Activities.cleanup_resources,
{ {
**metadata, **metadata,
'run_dir': run_dir,
'bucket_name': bucket_name, 'bucket_name': bucket_name,
'file_name': file_name, 'file_name': file_name,
}, },

7
pyproject.toml
View File

@@ -113,6 +113,10 @@ ignore_missing_imports = true
module = "sklearn.*" module = "sklearn.*"
ignore_missing_imports = true ignore_missing_imports = true
[[tool.mypy.overrides]]
module = "sientia_model.*"
ignore_missing_imports = true
[[tool.mypy.overrides]] [[tool.mypy.overrides]]
module = "yaml" module = "yaml"
ignore_missing_imports = true ignore_missing_imports = true
@@ -136,6 +140,9 @@ markers = [
"unit: marks tests as unit tests", "unit: marks tests as unit tests",
] ]
[tool.pyright]
reportMissingTypeStubs = false
[tool.coverage.run] [tool.coverage.run]
source = ["model_manager"] source = ["model_manager"]
omit = [ omit = [

7
requirements.txt
View File

@@ -3,10 +3,7 @@ psycopg2-binary==2.9.11
sqlalchemy==2.0.44 sqlalchemy==2.0.44
boto3==1.40.55 boto3==1.40.55
botocore==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 prometheus-client==0.23.1
mlflow==2.18.0
setuptools<81
evidently==0.4.21
beautifulsoup4==4.12.3 beautifulsoup4==4.12.3
scikit-learn==1.4.2

23
scripts/run_training_test.py
View File

@@ -289,17 +289,18 @@ def run_local_pipeline(
csv_path: Path, csv_path: Path,
save_mlflow: bool = False, save_mlflow: bool = False,
) -> dict: ) -> 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. Reads CSV from disk, runs DataManagerRepository.prepare_training_data and
Optionally saves to MLflow if save_mlflow is True (requires MLflow env). compute_regression_metrics. Optionally saves to MLflow if save_mlflow is
True (requires MLflow env).
Returns: Returns:
dict: {'success': bool, 'error': str | None, 'scenario': str, ...} dict: {'success': bool, 'error': str | None, 'scenario': str, ...}
""" """
from model_manager.utils.logger_helper import get_logger from model_manager.utils.logger_helper import get_logger
from model_manager.utils.models.train_model_params import TrainModelParams 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 = { result = {
'scenario': scenario_name, 'scenario': scenario_name,
@@ -327,14 +328,22 @@ def run_local_pipeline(
return result return result
logger = get_logger(__name__) logger = get_logger(__name__)
training_repository = TrainingRepository(logger) data_manager_repository = DataManagerRepository(logger)
with open(csv_path, 'rb') as f: with open(csv_path, 'rb') as f:
file_content = BytesIO(f.read()) file_content = BytesIO(f.read())
try: try:
train_result = training_repository.train(file_content, train_params) train_result = data_manager_repository.prepare_training_data(
train_result = training_repository.after_train_calculation(train_params, train_result) 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: except Exception as e:
result['error'] = str(e) result['error'] = str(e)
raise # re-raise so caller gets full traceback for diagnosis raise # re-raise so caller gets full traceback for diagnosis

316
tests/activities/test_activities.py
View File

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

11
tests/utils/models/test_train_model_result.py
View File

@@ -61,7 +61,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
x_train, x_test, y_train, y_test = sample_dataframes x_train, x_test, y_train, y_test = sample_dataframes
process_data = MagicMock() process_data = MagicMock()
regr = MagicMock() regr = MagicMock()
scaler_dict = {'var1': {'min': 0, 'max': 100}}
result = TrainModelResult( result = TrainModelResult(
params=sample_params, params=sample_params,
@@ -71,7 +70,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes):
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=regr, regr=regr,
scaler_dict=scaler_dict,
) )
assert result.params == sample_params 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_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
scaler_dict={},
) )
assert result.y_pred is None 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_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
scaler_dict={},
y_pred=y_pred, y_pred=y_pred,
mse_val=1.5, mse_val=1.5,
mae_val=1.2, 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_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
scaler_dict={},
run_name='test-experiment-1', run_name='test-experiment-1',
report_path='/path/to/report.html', report_path='/path/to/report.html',
train_data_path='/path/to/train_data.csv', 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(): 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 from dataclasses import fields
result_fields = fields(TrainModelResult) result_fields = fields(TrainModelResult)
assert len(result_fields) == 20 assert len(result_fields) == 19
field_names = {f.name for f in result_fields} field_names = {f.name for f in result_fields}
expected_fields = { expected_fields = {
@@ -201,7 +196,6 @@ def test_train_model_result_field_count():
'y_train', 'y_train',
'y_test', 'y_test',
'regr', 'regr',
'scaler_dict',
'y_pred', 'y_pred',
'y_train_pred', 'y_train_pred',
'mse_val', 'mse_val',
@@ -231,7 +225,6 @@ def test_train_model_result_complete_workflow(sample_params, sample_dataframes):
y_train=y_train, y_train=y_train,
y_test=y_test, y_test=y_test,
regr=MagicMock(), regr=MagicMock(),
scaler_dict={'var1': {'min': 0, 'max': 100}},
) )
# Step 2: Add predictions and metrics # Step 2: Add predictions and metrics

982
tests/utils/repository/test_model_repository.py
View File

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

513
tests/utils/repository/test_storage_repository.py
View File

@@ -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')

1329
tests/utils/repository/test_training_repository.py
View File

File diff suppressed because it is too large Load Diff

196
tests/worker/test_worker.py
View File

@@ -20,6 +20,10 @@ def mock_env_vars():
'PROJECT_NAME': 'test-project', 'PROJECT_NAME': 'test-project',
'TRAIN_TASK_QUEUE': 'train_model-local_queue', 'TRAIN_TASK_QUEUE': 'train_model-local_queue',
'CLEANUP_TASK_QUEUE': 'cleanup-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): 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.Runtime')
@patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.Activities')
@patch('model_manager.worker.worker.NotificationHandler') @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_mongodb_config')
@patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_postgres_config')
@patch('model_manager.worker.worker.build_mlflow_config') @patch('model_manager.worker.worker.build_mlflow_config')
@@ -191,6 +197,8 @@ async def test_main_successful_startup(
mock_build_mlflow, mock_build_mlflow,
mock_build_postgres, mock_build_postgres,
mock_build_mongodb, mock_build_mongodb,
mock_build_plugin_store_config,
mock_plugin_store_class,
mock_notification_handler_class, mock_notification_handler_class,
mock_activities_class, mock_activities_class,
mock_runtime_class, mock_runtime_class,
@@ -220,6 +228,23 @@ async def test_main_successful_startup(
mock_notification_handler_class.return_value = mock_notification_handler mock_notification_handler_class.return_value = mock_notification_handler
mock_activities_class.return_value = mock_activities 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 = Mock()
mock_runtime_class.return_value = mock_runtime 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.Runtime')
@patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.Activities')
@patch('model_manager.worker.worker.NotificationHandler') @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_mongodb_config')
@patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_postgres_config')
@patch('model_manager.worker.worker.build_mlflow_config') @patch('model_manager.worker.worker.build_mlflow_config')
@@ -277,6 +304,8 @@ async def test_main_handles_exception(
mock_build_mlflow, mock_build_mlflow,
mock_build_postgres, mock_build_postgres,
mock_build_mongodb, mock_build_mongodb,
mock_build_plugin_store_config,
mock_plugin_store_class,
mock_notification_handler_class, mock_notification_handler_class,
mock_activities_class, mock_activities_class,
mock_runtime_class, mock_runtime_class,
@@ -320,6 +349,23 @@ async def test_main_handles_exception(
mock_app_up = Mock() mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up 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 # Run main() and expect SystemExit
with pytest.raises(SystemExit) as exc_info: with pytest.raises(SystemExit) as exc_info:
await main() await main()
@@ -343,6 +389,8 @@ async def test_main_handles_exception(
@patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.Activities')
@patch('model_manager.worker.worker.NotificationHandler') @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_mongodb_config')
@patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_postgres_config')
@patch('model_manager.worker.worker.build_mlflow_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_mlflow,
mock_build_postgres, mock_build_postgres,
mock_build_mongodb, mock_build_mongodb,
mock_build_plugin_store_config,
mock_plugin_store_class,
mock_notification_handler_class, mock_notification_handler_class,
mock_activities_class, mock_activities_class,
mock_runtime_class, mock_runtime_class,
@@ -377,6 +427,10 @@ async def test_main_temporal_client_configuration(
'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_HOST': 'temporal.example.com:7233',
'TEMPORAL_NAMESPACE': 'production', 'TEMPORAL_NAMESPACE': 'production',
'TEMPORAL_USE_TLS': 'true', '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 # Setup mocks
@@ -390,6 +444,23 @@ async def test_main_temporal_client_configuration(
mock_build_mlflow.return_value = {} mock_build_mlflow.return_value = {}
mock_build_minio.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 = Mock()
mock_notification_handler.shutdown = Mock() mock_notification_handler.shutdown = Mock()
mock_notification_handler_class.return_value = mock_notification_handler 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.Runtime')
@patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.Activities')
@patch('model_manager.worker.worker.NotificationHandler') @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_mongodb_config')
@patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_postgres_config')
@patch('model_manager.worker.worker.build_mlflow_config') @patch('model_manager.worker.worker.build_mlflow_config')
@@ -446,6 +519,8 @@ async def test_main_worker_configuration(
mock_build_mlflow, mock_build_mlflow,
mock_build_postgres, mock_build_postgres,
mock_build_mongodb, mock_build_mongodb,
mock_build_plugin_store_config,
mock_plugin_store_class,
mock_notification_handler_class, mock_notification_handler_class,
mock_activities_class, mock_activities_class,
mock_runtime_class, mock_runtime_class,
@@ -500,6 +575,23 @@ async def test_main_worker_configuration(
mock_app_up = Mock() mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up 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() # Run main()
with pytest.raises(SystemExit): with pytest.raises(SystemExit):
await main() await main()
@@ -539,6 +631,8 @@ async def test_main_worker_configuration(
@patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Runtime')
@patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.Activities')
@patch('model_manager.worker.worker.NotificationHandler') @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_mongodb_config')
@patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_postgres_config')
@patch('model_manager.worker.worker.build_mlflow_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_mlflow,
mock_build_postgres, mock_build_postgres,
mock_build_mongodb, mock_build_mongodb,
mock_build_plugin_store_config,
mock_plugin_store_class,
mock_notification_handler_class, mock_notification_handler_class,
mock_activities_class, mock_activities_class,
mock_runtime_class, mock_runtime_class,
@@ -603,6 +699,106 @@ async def test_main_schedule_creation_failure_does_not_stop_worker(
mock_app_up = Mock() mock_app_up = Mock()
mock_metrics.APP_UP.labels.return_value = mock_app_up 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 # Run main() - should not fail despite schedule creation error
with pytest.raises(SystemExit): with pytest.raises(SystemExit):
await main() await main()

302
values.yaml
View File

@@ -1,58 +1,33 @@
# 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. # This is a YAML-formatted file.
# Declare variables to be passed into your templates. # 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/ projectName: &projectName "sientia-dataops-model-manager"
replicaCount: 1
# 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: image:
repository: aignosi.azurecr.io/sientia-dataops-model-manager repository: aignosi.azurecr.io/sientia-module
# This sets the pull policy for images.
pullPolicy: IfNotPresent pullPolicy: IfNotPresent
# Overrides the image tag whose default is the chart appVersion. tag: "1.0.0"
tag: "1.1.2"
# 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: # Global configuration shared by all runtimes
- name: docker-hub-secret # -----------------------------------------------------------------------------
global:
# Namespace used by the chart.
namespace: sientia
# This is to override the chart name. # Common labels applied to pods (can be extended per project).
nameOverride: "sientia-dataops-model-manager" commonLabels: {}
fullnameOverride: "sientia-dataops-model-manager"
namespace: sientia
# This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/ # Resources inherited by all runtimes unless overridden.
serviceAccount: resources: {}
# Specifies whether a service account should be created
create: true
# Automatically mount a ServiceAccount's API credentials?
automount: true
# Annotations to add to the service account
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"
# This is for setting Kubernetes Annotations to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
podAnnotations: {}
# This is for setting Kubernetes Labels to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/
podLabels: {}
podSecurityContext: {}
# fsGroup: 2000
securityContext: {}
# capabilities:
# drop:
# - ALL
# readOnlyRootFilesystem: true
# runAsNonRoot: true
# runAsUser: 1000
resources: {}
# We usually recommend not to specify default resources and to leave this as a conscious # 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 # 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 # resources, such as Minikube. If you do want to specify resources, uncomment the following
@@ -64,8 +39,10 @@ resources: {}
# cpu: 100m # cpu: 100m
# memory: 128Mi # 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/ # Probes inherited by all runtimes unless overridden.
livenessProbe: # More information:
# https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/
livenessProbe:
exec: exec:
command: command:
- python3 - python3
@@ -74,7 +51,7 @@ livenessProbe:
initialDelaySeconds: 20 initialDelaySeconds: 20
periodSeconds: 30 periodSeconds: 30
readinessProbe: readinessProbe:
exec: exec:
command: command:
- python3 - python3
@@ -83,75 +60,18 @@ readinessProbe:
initialDelaySeconds: 10 initialDelaySeconds: 10
periodSeconds: 15 periodSeconds: 15
# This section is for setting up autoscaling more information can be found here: https://kubernetes.io/docs/concepts/workloads/autoscaling/ # Autoscaling configuration inherited by all runtimes unless overridden.
autoscaling: # More information:
# https://kubernetes.io/docs/concepts/workloads/autoscaling/
autoscaling:
enabled: false enabled: false
minReplicas: 1 minReplicas: 1
maxReplicas: 100 maxReplicas: 100
targetCPUUtilizationPercentage: 80 targetCPUUtilizationPercentage: 80
# targetMemoryUtilizationPercentage: 80 # targetMemoryUtilizationPercentage: 80
# Additional volumes on the output Deployment definition. # Environment variables shared by all runtimes.
volumes: env:
- name: reports-volume
emptyDir:
sizeLimit: 1Gi
# Additional volumeMounts on the output Deployment definition.
volumeMounts:
- name: reports-volume
mountPath: "/app/model_manager/reports/temp"
# Deployment strategy configuration
# More information: https://kubernetes.io/docs/concepts/workloads/controllers/deployment/#strategy
deploymentStrategy:
type: Recreate
# rollingUpdate:
# maxSurge: 0
# maxUnavailable: 1
# Number of old ReplicaSets to retain
revisionHistoryLimit: 2
nodeSelector: {}
tolerations: []
affinity: {}
services:
sdk-metrics:
enabled: true
type: ClusterIP
port: 9091
targetPort: 9091
name: sdk-metrics
metrics:
enabled: true
type: ClusterIP
port: 9090
targetPort: 9090
name: metrics
# Configuração do ServiceMonitor para o Prometheus Operator
# ref: https://github.com/prometheus-operator/prometheus-operator
serviceMonitor:
# Se true, um recurso ServiceMonitor será criado.
enabled: true
# O intervalo no qual as métricas devem ser coletadas (ex: 30s, 1m).
endpoints:
- port: metrics
path: /metrics
interval: 30s
relabelings: []
- port: sdk-metrics
path: /metrics
interval: 30s
relabelings: []
additionalLabels:
release: kube-prometheus-stack
env:
- name: POSTGRES_HOST - name: POSTGRES_HOST
value: "paradedb-rw.paradedb.svc.cluster.local" value: "paradedb-rw.paradedb.svc.cluster.local"
- name: POSTGRES_PORT - name: POSTGRES_PORT
@@ -180,8 +100,6 @@ env:
value: "9090" value: "9090"
- name: HTTP_SDK_METRICS_PORT - name: HTTP_SDK_METRICS_PORT
value: "9091" value: "9091"
- name: PROJECT_NAME
value: "sientia-model-manager"
- name: TEMPORAL_HOST - name: TEMPORAL_HOST
value: "temporal-frontend.temporal.svc.cluster.local:7233" value: "temporal-frontend.temporal.svc.cluster.local:7233"
@@ -194,6 +112,25 @@ env:
- name: TEMPORAL_USE_TLS - name: TEMPORAL_USE_TLS
value: "false" 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 - name: MONGODB_USERNAME
value: "root" value: "root"
- name: MONGODB_PASSWORD - name: MONGODB_PASSWORD
@@ -256,13 +193,120 @@ env:
- name: CLEANUP_EXECUTION_TIMEOUT_HOURS - name: CLEANUP_EXECUTION_TIMEOUT_HOURS
value: "1" value: "1"
- name: EXTRA_PIP_REQUIREMENTS - name: PYPI_SERVER
value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git" value: "http://library-distribution-server.library.svc.cluster.local:5000"
- name: POD_ID # -----------------------------------------------------------------------------
valueFrom: # Runtimes configuration
fieldRef: # -----------------------------------------------------------------------------
fieldPath: metadata.name # 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
# This is to override the chart name.
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:
# Specifies whether a service account should be created
create: true
# Automatically mount a ServiceAccount's API credentials?
automount: true
# Annotations to add to the service account
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: *projectName
# This is for setting Kubernetes Annotations to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/
podAnnotations: {}
# This is for setting Kubernetes Labels to a Pod.
# For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/
podLabels: {}
podSecurityContext: {}
# fsGroup: 2000
securityContext: {}
# capabilities:
# drop:
# - ALL
# readOnlyRootFilesystem: true
# runAsNonRoot: true
# runAsUser: 1000
# Additional volumes on the output Deployment definition.
volumes:
- name: reports-volume
emptyDir:
sizeLimit: 1Gi
# Additional volumeMounts on the output Deployment definition.
volumeMounts:
- name: reports-volume
mountPath: "/app/model_manager/reports/temp"
# Deployment strategy configuration
# More information: https://kubernetes.io/docs/concepts/workloads/controllers/deployment/#strategy
deploymentStrategy:
type: Recreate
# rollingUpdate:
# maxSurge: 0
# maxUnavailable: 1
# Number of old ReplicaSets to retain
revisionHistoryLimit: 2
nodeSelector: {}
tolerations: []
affinity: {}
services:
sdk-metrics:
enabled: true
type: ClusterIP
port: 9091
targetPort: 9091
name: sdk-metrics
metrics:
enabled: true
type: ClusterIP
port: 9090
targetPort: 9090
name: metrics
# Configuração do ServiceMonitor para o Prometheus Operator
# ref: https://github.com/prometheus-operator/prometheus-operator
serviceMonitor:
# Se true, um recurso ServiceMonitor será criado.
enabled: true
# O intervalo no qual as métricas devem ser coletadas (ex: 30s, 1m).
endpoints:
- port: metrics
path: /metrics
interval: 30s
relabelings: []
- port: sdk-metrics
path: /metrics
interval: 30s
relabelings: []
additionalLabels:
release: kube-prometheus-stack
ssh: ssh:
enabled: false enabled: false
@@ -308,11 +352,21 @@ grafanaDatasource:
# jsonData: # jsonData:
# timeInterval: "5s" # 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 usage examples
# helm upgrade --install sientia-dataops-model-manager sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.0 # -----------------------------------------------------------------------------
# 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 \ # kubectl create secret generic git-ssh-key-sientia-model-manager-worker \
# --namespace sientia \ # --namespace sientia \
# --from-file=ssh-privatekey=git_key \ # --from-file=ssh-privatekey=git_key \
# --type=kubernetes.io/ssh-auth # --type=kubernetes.io/ssh-auth