diff --git a/model-manager-plugin-store-migration-plan.md b/model-manager-plugin-store-migration-plan.md new file mode 100644 index 0000000..63456f0 --- /dev/null +++ b/model-manager-plugin-store-migration-plan.md @@ -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]] + diff --git a/model_manager/activities/activities.py b/model_manager/activities/activities.py index 27ad6da..66d1bdd 100644 --- a/model_manager/activities/activities.py +++ b/model_manager/activities/activities.py @@ -6,12 +6,14 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository + from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository + from sientia_model.model_repository.plugin_store import PluginStore from model_manager.activities.cleanup import Cleanup from model_manager.activities.experiment_tracking import ExperimentTracking from model_manager.activities.training import Training - from model_manager.utils.repository.model_repository import ModelRepository - from model_manager.utils.repository.storage_repository import StorageRepository class Activities(ExperimentTracking, Training, Cleanup): @@ -40,8 +42,10 @@ class Activities(ExperimentTracking, Training, Cleanup): postgres_config: dict[str, Any], mlflow_config: dict[str, Any], minio_config: dict[str, Any], + plugin_store: PluginStore, logger: Logger, notification_handler: NotificationHandler, + metrics_controller: MetricsController, ): """ Initialize the Activities orchestrator with all required configurations. @@ -62,9 +66,7 @@ class Activities(ExperimentTracking, Training, Cleanup): Raises: Exception: If any parent class initialization fails """ - metrics_controller = MetricsController( - logger=logger, - ) + ExperimentTracking.__init__( self, @@ -80,30 +82,40 @@ class Activities(ExperimentTracking, Training, Cleanup): metrics_controller=metrics_controller, ) - self.model_repository = ModelRepository( - url=mlflow_config['url'], + self.mlflow_repository = SientiaMLflowRepository( + host=mlflow_config['url'], username=mlflow_config['username'], password=mlflow_config['password'], logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, ) - self.storage_repository = StorageRepository( - endpoint_url=minio_config['endpoint_url'], + # MinIO repository used for all object storage operations + endpoint_url = minio_config['endpoint_url'] + # MinioRepository expects the endpoint without scheme + if endpoint_url.startswith('http://'): + endpoint = endpoint_url.removeprefix('http://') + elif endpoint_url.startswith('https://'): + endpoint = endpoint_url.removeprefix('https://') + else: + endpoint = endpoint_url + + self.minio_repository = MinioRepository( + endpoint=endpoint, access_key=minio_config['access_key'], secret_key=minio_config['secret_key'], - region=minio_config['region'], - use_ssl=minio_config['use_ssl'], - max_retry_attempts=minio_config['max_retry_attempts'], - retry_mode=minio_config['retry_mode'], - connect_timeout=minio_config['connect_timeout'], - read_timeout=minio_config['read_timeout'], logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + secure=minio_config['use_ssl'], ) Training.__init__( self, - model_repository=self.model_repository, - storage_repository=self.storage_repository, + mlflow_repository=self.mlflow_repository, + plugin_store=plugin_store, + minio_repository=self.minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -111,7 +123,7 @@ class Activities(ExperimentTracking, Training, Cleanup): Cleanup.__init__( self, - storage_repository=self.storage_repository, + minio_repository=self.minio_repository, logger=logger, notification_handler=notification_handler, metrics_controller=metrics_controller, @@ -137,7 +149,7 @@ class Activities(ExperimentTracking, Training, Cleanup): # Logging here could cause issues if logger is already destroyed pass - async def shutdown(self): + def shutdown(self): """ Gracefully shutdown all activities and clean up resources. @@ -152,4 +164,4 @@ class Activities(ExperimentTracking, Training, Cleanup): """ ExperimentTracking.close(self) self.info('Postgres client closed') - self.storage_repository.close() + SientiaMonitoring.shutdown(self) diff --git a/model_manager/activities/cleanup.py b/model_manager/activities/cleanup.py index 1f9654e..fcf22b0 100644 --- a/model_manager/activities/cleanup.py +++ b/model_manager/activities/cleanup.py @@ -21,9 +21,9 @@ with workflow.unsafe.imports_passed_through(): from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL - from model_manager.utils.repository.storage_repository import StorageRepository RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24')) DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true' @@ -41,7 +41,7 @@ class Cleanup(SientiaMonitoring): def __init__( self, - storage_repository: StorageRepository, + minio_repository: MinioRepository, logger: Logger, notification_handler: NotificationHandler, metrics_controller: MetricsController, @@ -50,13 +50,13 @@ class Cleanup(SientiaMonitoring): Initialize Cleanup activity. Args: - storage_repository: Repository for MinIO operations + minio_repository: Repository for MinIO operations logger: Logger instance for observability notification_handler: Handler for sending notifications metrics_controller: Controller for metrics emission """ SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) - self.storage_repository = storage_repository + self.minio_repository = minio_repository # Configuration from environment variables self.retention_hours = RETENTION_HOURS @@ -113,15 +113,21 @@ class Cleanup(SientiaMonitoring): files_deleted = 0 errors = [] - # List objects in the specified bucket - max_keys = self.max_keys_cleanup # Use environment variable for page size - objects = self.storage_repository.list_bucket_objects(bucket_name, max_keys) + # List objects in the specified bucket. MinioRepository applies BASE_PREFIX + # internally; we request all objects under that prefix for this bucket. + objects = await self.minio_repository.list_objects( + prefix='', + bucket=bucket_name, + recursive=True, + metadata=metadata, + ) for obj_key in objects: files_scanned += 1 - # Extract timestamp from filename - match = self.minio_timestamp_pattern.match(obj_key) + # Extract timestamp from the filename portion of the object key + filename = obj_key.split('/')[-1] + match = self.minio_timestamp_pattern.match(filename) if not match: self.debug(f'Skipping file without timestamp pattern: {obj_key}', metadata) continue @@ -137,10 +143,14 @@ class Cleanup(SientiaMonitoring): files_deleted += 1 else: try: - self.storage_repository.delete_file(bucket_name, obj_key) + await self.minio_repository.delete_file( + object_name=obj_key, + bucket=None, + metadata=metadata, + ) self.info(f'Deleted stale file: {obj_key}', metadata) files_deleted += 1 - except OSError as e: + except Exception as e: # noqa: BLE001 error_msg = f'Failed to delete {obj_key}: {str(e)}' errors.append(error_msg) self.error(error_msg, metadata) diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index aaae6b2..88e16b7 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -12,18 +12,20 @@ with workflow.unsafe.imports_passed_through(): import traceback from typing import Any + import pandas as pd from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.notifications.models import NotificationLevel from sientia_do.observability.logger import Logger from sientia_do.observability.metrics_controller import MetricsController from sientia_do.observability.sientia_monitoring import SientiaMonitoring + from sientia_do.repository.minio_repository import MinioRepository + from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository + from sientia_model.model_repository.plugin_store import PluginStore from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL from model_manager.utils.exceptions import ModelTrainingError from model_manager.utils.models.train_model_params import TrainModelParams - from model_manager.utils.repository.model_repository import ModelRepository - from model_manager.utils.repository.storage_repository import StorageRepository - from model_manager.utils.repository.training_repository import TrainingRepository + from model_manager.utils.repository.data_manager_repository import DataManagerRepository class Training(SientiaMonitoring): @@ -38,8 +40,9 @@ class Training(SientiaMonitoring): def __init__( self, - model_repository: ModelRepository, - storage_repository: StorageRepository, + mlflow_repository: SientiaMLflowRepository, + plugin_store: PluginStore, + minio_repository: MinioRepository, logger: Logger, notification_handler: NotificationHandler, metrics_controller: MetricsController, @@ -52,9 +55,10 @@ class Training(SientiaMonitoring): notification_handler: Handler for sending notifications """ SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller) - self.training_repository = TrainingRepository(logger) - self.model_repository = model_repository - self.storage_repository = storage_repository + self.data_manager_repository = DataManagerRepository(logger) + self.mlflow_repository = mlflow_repository + self.plugin_store = plugin_store + self.minio_repository = minio_repository @activity.defn(name='validate_train_params') async def validate_train_params(self, input_data: dict[str, Any]) -> TrainModelParams: @@ -120,8 +124,8 @@ class Training(SientiaMonitoring): This activity orchestrates the ML training pipeline: 1. Validate input parameters. - 2. Train the model via TrainingRepository. - 3. Perform post-training calculations. + 2. Prepare data via DataManagerRepository. + 3. Train the model and compute metrics. Args: input_data: Training configuration containing: @@ -130,41 +134,98 @@ class Training(SientiaMonitoring): - train_params (TrainModelParams | dict): Training parameters. Returns: - dict: Keys `run_name` and `run_dir` when training and saving succeed. + dict: Key `run_name` when training and saving succeed. Raises: ValueError: If input validation fails. Exception: If training fails (after sending notification). """ - metadata = input_data.get('metadata', {}) + metadata = input_data.get('metadata') train_params = input_data['train_params'] if isinstance(train_params, dict): train_params = TrainModelParams.from_dict(train_params) - # type: ignore[assignment] + model_trained = False model_saved = False metrics_status = 'success' try: - with self.storage_repository.fetch_file( - train_params.bucket_name, train_params.file_name - ) as uploaded_file: - train_result = self.training_repository.train(uploaded_file, train_params) + # Download training file bytes from MinIO + train_bytes = await self.minio_repository.download_file( + object_name=train_params.file_name, + bucket=train_params.bucket_name, + metadata=metadata, + ) - train_result = self.training_repository.after_train_calculation( - train_params, train_result + # Download optional validation file bytes from the same bucket + val_bytes: bytes | None = None + validation_name = getattr(train_params, 'validation_file_name', None) + if validation_name is not None: + val_bytes = await self.minio_repository.download_file( + object_name=validation_name, + bucket=train_params.bucket_name, + metadata=metadata, ) - model_trained = True - train_result = self.model_repository.save_model(train_result) - model_saved = True + train_result = self.data_manager_repository.prepare_training_data( + train_file_bytes=train_bytes, + validation_file_bytes=val_bytes, + params=train_params, + metadata=metadata, + ) - return { - 'run_name': train_result.run_name, - 'run_dir': train_result.run_dir, - } + wrapper = await self.plugin_store.get_model( + model_name=train_params.model_name, + force_download=False, + opt_params={}, + model_kwargs={}, + data_model_kwargs={}, + metadata=metadata, + ) + + train_df = pd.concat([train_result.x_train, train_result.y_train], axis=1) + val_df = pd.concat([train_result.x_test, train_result.y_test], axis=1) + + wrapper.train( + train_data=train_df, + val_data=val_df, + target=train_params.target_variable, + ) + + # Generate predictions using the trained wrapper + transformed_train, _ = wrapper.transform(train_df) + transformed_val, _ = wrapper.transform(val_df) + + y_train_pred_df, _ = wrapper.predict({}, transformed_train) + y_val_pred_df, _ = wrapper.predict({}, transformed_val) + + # Use the first column of the prediction DataFrame as the target prediction + train_result.y_train_pred = y_train_pred_df.iloc[:, 0] + train_result.y_pred = y_val_pred_df.iloc[:, 0] + + train_result = self.data_manager_repository.compute_regression_metrics( + train_params, + train_result, + ) + + model_trained = True + + async with self.mlflow_repository.start_run( + model_name=train_params.model_name, + run_name=None, + experiment_name=train_params.experiment_name, + tags=None, + metadata=metadata, + ) as run_info: + wrapper.store_model(name=train_params.model_name) + + model_saved = True + + return { + 'run_name': run_info.run_name or run_info.run_id, + } except Exception as e: # noqa: BLE001 metrics_status = 'error' @@ -205,7 +266,6 @@ class Training(SientiaMonitoring): Args: input_data: Cleanup configuration containing: - metadata (dict): Workflow execution metadata. - - run_dir (str): Temporary directory to remove. - bucket_name (str): MinIO bucket of the uploaded file. - file_name (str): MinIO object key to delete. @@ -213,19 +273,21 @@ class Training(SientiaMonitoring): Exception: If cleanup fails (after sending notification). """ metadata = input_data.get('metadata', {}) - run_dir = input_data.get('run_dir', '') bucket_name = input_data.get('bucket_name', '') file_name = input_data.get('file_name', '') metrics_status = 'success' try: - self.model_repository.cleanup_run_directory(run_dir) - self.storage_repository.delete_file(bucket_name, file_name) + await self.minio_repository.delete_file( + object_name=file_name, + bucket=bucket_name, + metadata=metadata, + ) except Exception as e: # noqa: BLE001 metrics_status = 'error' error_msg = ( - f'Error cleaning up resources - Run directory: {run_dir}, ' + 'Error cleaning up resources - ' f'File: {bucket_name}/{file_name}, Error: {str(e)}' ) diff --git a/model_manager/utils/connectors_config.py b/model_manager/utils/connectors_config.py index 9a1e40a..61853f0 100644 --- a/model_manager/utils/connectors_config.py +++ b/model_manager/utils/connectors_config.py @@ -121,3 +121,43 @@ def build_minio_config() -> dict[str, Any]: 'connect_timeout': int(getenv('MINIO_CONNECT_TIMEOUT', '10')), 'read_timeout': int(getenv('MINIO_READ_TIMEOUT', '60')), } + + +def build_plugin_store_config() -> dict[str, Any]: + """ + Build PluginStore configuration from environment variables. + + This function constructs a configuration dictionary for the PluginStore + client using environment variables with sensible defaults for local + development. + + Environment Variables: + STORE_BASE_URL: Base URL of the PluginStore backing Git server + (default: http://localhost:3000) + STORE_OWNER: Repository owner/organization (default: sientia) + STORE_REPO: Repository name (default: model-library-store) + STORE_BRANCH: Optional branch name + STORE_USERNAME: Optional username for Git HTTP authentication + STORE_PASSWORD: Optional password/token for Git HTTP authentication + PYPI_SERVER: Optional custom PyPI index URL for runtime installation + (default: http://localhost:5000) + PYPI_USERNAME: Optional username for PyPI authentication + PYPI_PASSWORD: Optional password/token for PyPI authentication + + Returns: + dict: PluginStore configuration dictionary with all connector parameters + """ + + cache_ttl_seconds = getenv('STORE_CACHE_TTL_SECONDS') + return { + 'base_url': getenv('STORE_BASE_URL', 'http://localhost:3000'), + 'owner': getenv('STORE_OWNER', 'sientia'), + 'repo': getenv('STORE_REPO', 'model-library-store'), + 'username': getenv('STORE_USERNAME'), + 'password': getenv('STORE_PASSWORD'), + 'branch': getenv('STORE_BRANCH'), + 'cache_ttl_seconds': int(cache_ttl_seconds) if cache_ttl_seconds else None, + 'pypi_index_url': getenv('PYPI_SERVER', 'http://localhost:5000'), + 'pypi_username': getenv('PYPI_USERNAME'), + 'pypi_password': getenv('PYPI_PASSWORD'), + } diff --git a/model_manager/utils/models/train_model_params.py b/model_manager/utils/models/train_model_params.py index 87a182d..1eea000 100644 --- a/model_manager/utils/models/train_model_params.py +++ b/model_manager/utils/models/train_model_params.py @@ -33,7 +33,8 @@ class TrainModelParams: use_scaler (bool): Whether to use a scaler for data normalization. include_ar (bool): Whether to include autoregressive variables. bucket_name (str): Name of the MinIO bucket containing training data. - file_name (str): Name of the file in the MinIO bucket. + file_name (str): Name of the training file in the MinIO bucket. + validation_file_name (str | None): Optional name of the validation file in the same MinIO bucket as the training file. line_separator (str): Line separator used in the CSV file. decimal_separator (str): Decimal separator used in the CSV file. date_column (str | None): Name of the date/time column. If set with date_format, the column is parsed as datetime. @@ -66,6 +67,7 @@ class TrainModelParams: include_ar: bool bucket_name: str file_name: str + validation_file_name: str | None line_separator: str decimal_separator: str date_column: str | None @@ -122,6 +124,9 @@ class TrainModelParams: include_ar=cls._check_none(data.get('include_ar'), bool, 'include_ar'), bucket_name=cls._check_none(data.get('bucket_name'), str, 'bucket_name'), file_name=cls._check_none(data.get('file_name'), str, 'file_name'), + validation_file_name=cls._check_type( + data.get('validation_file_name'), str, 'validation_file_name' + ), line_separator=cls._check_none(data.get('line_separator'), str, 'line_separator'), decimal_separator=cls._check_none( data.get('decimal_separator'), str, 'decimal_separator' @@ -220,6 +225,7 @@ class TrainModelParams: self._validate_limits() self._validate_required_strings() self._validate_date_format() + self._validate_validation_file_name() def _validate_numeric_ranges(self) -> None: """Validate numeric parameters are within acceptable ranges.""" @@ -335,3 +341,21 @@ class TrainModelParams: """Validate date_format is one of the allowed frontend formats when set.""" if self.date_format: validate_frontend_date_format(self.date_format) + + def _validate_validation_file_name(self) -> None: + """ + Validate that validation_file_name, when provided, is not empty or whitespace. + + This field is optional; when present it must point to a valid object key in the + same MinIO bucket specified by bucket_name. + """ + if self.validation_file_name is None: + return + + if not isinstance(self.validation_file_name, str): + raise TypeError( + f'validation_file_name must be a string, got {type(self.validation_file_name).__name__}' + ) + + if not self.validation_file_name.strip(): + raise ValueError('validation_file_name cannot be empty or whitespace') diff --git a/model_manager/utils/models/train_model_result.py b/model_manager/utils/models/train_model_result.py index 049bcd2..9b8d9cd 100644 --- a/model_manager/utils/models/train_model_result.py +++ b/model_manager/utils/models/train_model_result.py @@ -2,7 +2,6 @@ from dataclasses import dataclass import pandas as pd -from model_manager.sientia.models import DataPreprocessor, LinearRegressionModel from model_manager.utils.models.train_model_params import TrainModelParams @@ -12,18 +11,15 @@ class TrainModelResult: A data container for storing the results of a machine learning training process. This dataclass encapsulates all outputs from the training pipeline, including - the trained model, datasets, evaluation metrics, and paths to generated artifacts. + the prepared datasets, evaluation metrics, and paths to generated artifacts. It is used to pass results between activities in the training workflow. Attributes: params (TrainModelParams): The parameters used to train the model. - process_data (DataPreprocessor): The data preprocessor object used to process the input data. x_train (pd.DataFrame): The training dataset features. x_test (pd.DataFrame): The testing dataset features. y_train (pd.DataFrame): The training dataset target values. y_test (pd.DataFrame): The testing dataset target values. - regr (LinearRegressionModel): The trained linear regression model. - scaler_dict (dict): A dictionary containing the scalers used to scale the features and target values. y_pred (pd.Series | None): The predicted target values for the testing dataset. Default is None. y_train_pred (pd.Series | None): The predicted target values for the training dataset. Default is None. mse_val (float | None): The Mean Squared Error (MSE) of the predictions. Default is None. @@ -35,17 +31,13 @@ class TrainModelResult: report_path (str | None): The path to the generated HTML report file. Default is None. train_data_path (str | None): The path to the training dataset CSV file. Default is None. test_data_path (str | None): The path to the testing dataset CSV file. Default is None. - run_dir (str | None): The path to the run directory containing all artifacts. Default is None. """ params: TrainModelParams - process_data: DataPreprocessor x_train: pd.DataFrame x_test: pd.DataFrame y_train: pd.Series y_test: pd.Series - regr: LinearRegressionModel - scaler_dict: dict y_pred: pd.Series | None = None y_train_pred: pd.Series | None = None mse_val: float | None = None @@ -57,4 +49,3 @@ class TrainModelResult: report_path: str | None = None train_data_path: str | None = None test_data_path: str | None = None - run_dir: str | None = None diff --git a/model_manager/utils/repository/data_manager_repository.py b/model_manager/utils/repository/data_manager_repository.py new file mode 100644 index 0000000..5a3ffc0 --- /dev/null +++ b/model_manager/utils/repository/data_manager_repository.py @@ -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 + diff --git a/model_manager/utils/repository/model_repository.py b/model_manager/utils/repository/model_repository.py deleted file mode 100644 index dcd1ba8..0000000 --- a/model_manager/utils/repository/model_repository.py +++ /dev/null @@ -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 "-". - """ - 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 diff --git a/model_manager/utils/repository/storage_repository.py b/model_manager/utils/repository/storage_repository.py deleted file mode 100644 index c8639fb..0000000 --- a/model_manager/utils/repository/storage_repository.py +++ /dev/null @@ -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 diff --git a/model_manager/utils/repository/training_repository.py b/model_manager/utils/repository/training_repository.py deleted file mode 100644 index 6a28b3e..0000000 --- a/model_manager/utils/repository/training_repository.py +++ /dev/null @@ -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 diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index f5b03b4..fcb60c6 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -40,6 +40,9 @@ with workflow.unsafe.imports_passed_through(): from prometheus_client import start_http_server from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler from sientia_do.observability.logger import Logger as SientiaLogger + from sientia_model.model_repository.plugin_store import PluginStore + from sientia_do.observability.metrics_controller import MetricsController + from sientia_do.temporal.worker.prepare_worker import prepare_worker from model_manager import metrics from model_manager.activities.activities import Activities @@ -48,6 +51,7 @@ with workflow.unsafe.imports_passed_through(): build_minio_config, build_mlflow_config, build_mongodb_config, + build_plugin_store_config, build_postgres_config, ) from model_manager.utils.logger_helper import get_logger @@ -55,6 +59,7 @@ with workflow.unsafe.imports_passed_through(): from model_manager.workflows.train_model import TrainModel POD_ID = os.getenv('POD_ID') +RUNTIME = os.getenv('RUNTIME') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue') CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue') @@ -79,13 +84,19 @@ async def main(): Exception: Any unhandled exception during worker execution SystemExit: On graceful shutdown or error conditions """ + + if not RUNTIME: + raise ValueError('RUNTIME environment variable is required') + host = os.getenv('TEMPORAL_HOST', 'localhost:7233') use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true' logger = get_logger(__name__) metadata = { 'pod_id': POD_ID, + 'runtime': RUNTIME, } + start_prometheus_server(logger, metadata) mongo_config = build_mongodb_config() @@ -99,12 +110,41 @@ async def main(): logger.custom_info(f'MongoDB client initialized at {mongo_config["uri"]}', metadata) + logger.custom_info(f'Initializing metrics controller', metadata) + + metrics_controller = MetricsController( + logger=logger + ) + + logger.custom_info(f'Installing runtime {RUNTIME}', metadata) + + plugin_store_parameters = build_plugin_store_config() + plugin_store = PluginStore( + base_url=plugin_store_parameters['base_url'], + owner=plugin_store_parameters['owner'], + repo=plugin_store_parameters['repo'], + username=plugin_store_parameters['username'], + password=plugin_store_parameters['password'], + branch=plugin_store_parameters['branch'], + cache_ttl_seconds=plugin_store_parameters['cache_ttl_seconds'], + pypi_index_url=plugin_store_parameters['pypi_index_url'], + pypi_username=plugin_store_parameters['pypi_username'], + pypi_password=plugin_store_parameters['pypi_password'], + logger=logger, + notification_handler=notification_handler, + metrics_controller=metrics_controller, + ) + + await plugin_store.install_runtime(runtime_name=RUNTIME) + activities = Activities( postgres_config=build_postgres_config(), mlflow_config=build_mlflow_config(), minio_config=build_minio_config(), + plugin_store=plugin_store, logger=logger, notification_handler=notification_handler, + metrics_controller=metrics_controller, ) new_runtime = Runtime( @@ -133,44 +173,33 @@ async def main(): # The schedule can be created manually if needed workers = [ - Worker( - temporal_client, - task_queue=TRAIN_TASK_QUEUE, - workflows=[TrainModel], + prepare_worker( + main_workflow=TrainModel, + other_workflows=[], activities=[ activities.update_experiment_run, activities.validate_train_params, activities.train_model, activities.cleanup_resources, ], - max_concurrent_workflow_tasks=10, - max_concurrent_activities=10, - max_concurrent_local_activities=10, - max_cached_workflows=100, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + temporal_client=temporal_client, + logger=logger, ), - Worker( - temporal_client, - task_queue=CLEANUP_TASK_QUEUE, - workflows=[CleanupFiles], + prepare_worker( + main_workflow=CleanupFiles, + other_workflows=[], activities=[ activities.cleanup_minio_files, activities.cleanup_temp_directories, ], - max_concurrent_workflow_tasks=20, - max_concurrent_activities=20, - max_concurrent_local_activities=20, - max_cached_workflows=100, - workflow_task_poller_behavior=PollerBehaviorAutoscaling(), - activity_task_poller_behavior=PollerBehaviorAutoscaling(), + temporal_client=temporal_client, + logger=logger, ), ] - handlers = [] - - for w in workers: - handlers.append(w.run()) + handlers = [ + w.run() for w in workers + ] logger.custom_info('Model manager workers initialized', metadata) @@ -183,7 +212,7 @@ async def main(): finally: notification_handler.shutdown() logger.custom_info('MongoDB client closed', metadata) - await activities.shutdown() + activities.shutdown() # Exit with a non-zero status code to indicate failure to Kubernetes metrics.APP_UP.labels(pod_id=POD_ID).set(0) # Mark app as DOWN sys.exit(1) diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index da1df5f..f3c37a4 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -118,7 +118,6 @@ class TrainModel: await self._cleanup_resources( experiment_run_id=experiment_run_id, - run_dir=(train_result.get('run_dir') or ''), bucket_name=train_params.bucket_name, file_name=train_params.file_name, metadata=metadata, @@ -275,7 +274,6 @@ class TrainModel: async def _cleanup_resources( self, experiment_run_id: int, - run_dir: str, bucket_name: str, file_name: str, metadata: dict[str, Any], @@ -283,12 +281,11 @@ class TrainModel: """ Cleanup resources and delete file from MinIO. - This method removes the temporary run directory via activity and deletes - the training file from MinIO. On success, updates DB status to FILE_DELETED. + This method deletes the training file from MinIO. On success, updates + DB status to FILE_DELETED. On error, updates DB status to FILE_DELETE_ERROR. Args: - saved_result: TrainModelResult with run_dir and params information experiment_run_id: Validated experiment run ID metadata: Workflow execution metadata @@ -300,7 +297,6 @@ class TrainModel: Activities.cleanup_resources, { **metadata, - 'run_dir': run_dir, 'bucket_name': bucket_name, 'file_name': file_name, }, diff --git a/pyproject.toml b/pyproject.toml index 0da188c..ea5d4cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,10 @@ ignore_missing_imports = true module = "sklearn.*" ignore_missing_imports = true +[[tool.mypy.overrides]] +module = "sientia_model.*" +ignore_missing_imports = true + [[tool.mypy.overrides]] module = "yaml" ignore_missing_imports = true @@ -136,6 +140,9 @@ markers = [ "unit: marks tests as unit tests", ] +[tool.pyright] +reportMissingTypeStubs = false + [tool.coverage.run] source = ["model_manager"] omit = [ diff --git a/requirements.txt b/requirements.txt index 4a28cc2..847ded1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,10 +3,7 @@ psycopg2-binary==2.9.11 sqlalchemy==2.0.44 boto3==1.40.55 botocore==1.40.55 -git+ssh://git@github.com/Aignosi/sientia-dataops-library.git@1.6.1 +git+ssh://git@github.com/Aignosi/sientia-model-library.git@0.7.2 +git+https://github.com/Aignosi/sientia-dataops-library.git@v1.10.1 prometheus-client==0.23.1 -mlflow==2.18.0 -setuptools<81 -evidently==0.4.21 beautifulsoup4==4.12.3 -scikit-learn==1.4.2 diff --git a/scripts/run_training_test.py b/scripts/run_training_test.py index a0ecfe0..a767368 100644 --- a/scripts/run_training_test.py +++ b/scripts/run_training_test.py @@ -289,17 +289,18 @@ def run_local_pipeline( csv_path: Path, save_mlflow: bool = False, ) -> dict: - """Run the same training pipeline locally (validate + train + after_train). + """Run the same training pipeline locally (validate + train + metrics). - Reads CSV from disk, runs TrainingRepository.train and after_train_calculation. - Optionally saves to MLflow if save_mlflow is True (requires MLflow env). + Reads CSV from disk, runs DataManagerRepository.prepare_training_data and + compute_regression_metrics. Optionally saves to MLflow if save_mlflow is + True (requires MLflow env). Returns: dict: {'success': bool, 'error': str | None, 'scenario': str, ...} """ from model_manager.utils.logger_helper import get_logger from model_manager.utils.models.train_model_params import TrainModelParams - from model_manager.utils.repository.training_repository import TrainingRepository + from model_manager.utils.repository.data_manager_repository import DataManagerRepository result = { 'scenario': scenario_name, @@ -327,14 +328,22 @@ def run_local_pipeline( return result logger = get_logger(__name__) - training_repository = TrainingRepository(logger) + data_manager_repository = DataManagerRepository(logger) with open(csv_path, 'rb') as f: file_content = BytesIO(f.read()) try: - train_result = training_repository.train(file_content, train_params) - train_result = training_repository.after_train_calculation(train_params, train_result) + train_result = data_manager_repository.prepare_training_data( + train_file_bytes=file_content.getvalue(), + validation_file_bytes=None, + params=train_params, + metadata={'source': 'run_local_pipeline', 'scenario': scenario_name}, + ) + train_result = data_manager_repository.compute_regression_metrics( + train_params, + train_result, + ) except Exception as e: result['error'] = str(e) raise # re-raise so caller gets full traceback for diagnosis diff --git a/tests/activities/test_activities.py b/tests/activities/test_activities.py deleted file mode 100644 index 4df4309..0000000 --- a/tests/activities/test_activities.py +++ /dev/null @@ -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 diff --git a/tests/utils/models/test_train_model_result.py b/tests/utils/models/test_train_model_result.py index 87adb2a..aa03df1 100644 --- a/tests/utils/models/test_train_model_result.py +++ b/tests/utils/models/test_train_model_result.py @@ -61,7 +61,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes): x_train, x_test, y_train, y_test = sample_dataframes process_data = MagicMock() regr = MagicMock() - scaler_dict = {'var1': {'min': 0, 'max': 100}} result = TrainModelResult( params=sample_params, @@ -71,7 +70,6 @@ def test_train_model_result_creation(sample_params, sample_dataframes): y_train=y_train, y_test=y_test, regr=regr, - scaler_dict=scaler_dict, ) assert result.params == sample_params @@ -96,7 +94,6 @@ def test_train_model_result_optional_fields_default_none(sample_params, sample_d y_train=y_train, y_test=y_test, regr=MagicMock(), - scaler_dict={}, ) assert result.y_pred is None @@ -123,7 +120,6 @@ def test_train_model_result_with_metrics(sample_params, sample_dataframes): y_train=y_train, y_test=y_test, regr=MagicMock(), - scaler_dict={}, y_pred=y_pred, mse_val=1.5, mae_val=1.2, @@ -148,7 +144,6 @@ def test_train_model_result_with_artifact_paths(sample_params, sample_dataframes y_train=y_train, y_test=y_test, regr=MagicMock(), - scaler_dict={}, run_name='test-experiment-1', report_path='/path/to/report.html', train_data_path='/path/to/train_data.csv', @@ -186,11 +181,11 @@ def test_train_model_result_is_dataclass(sample_params, sample_dataframes): def test_train_model_result_field_count(): - """Test that TrainModelResult has exactly 20 fields.""" + """Test that TrainModelResult has exactly 19 fields.""" from dataclasses import fields result_fields = fields(TrainModelResult) - assert len(result_fields) == 20 + assert len(result_fields) == 19 field_names = {f.name for f in result_fields} expected_fields = { @@ -201,7 +196,6 @@ def test_train_model_result_field_count(): 'y_train', 'y_test', 'regr', - 'scaler_dict', 'y_pred', 'y_train_pred', 'mse_val', @@ -231,7 +225,6 @@ def test_train_model_result_complete_workflow(sample_params, sample_dataframes): y_train=y_train, y_test=y_test, regr=MagicMock(), - scaler_dict={'var1': {'min': 0, 'max': 100}}, ) # Step 2: Add predictions and metrics diff --git a/tests/utils/repository/test_model_repository.py b/tests/utils/repository/test_model_repository.py deleted file mode 100644 index c7d6b63..0000000 --- a/tests/utils/repository/test_model_repository.py +++ /dev/null @@ -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 diff --git a/tests/utils/repository/test_storage_repository.py b/tests/utils/repository/test_storage_repository.py deleted file mode 100644 index 3aaeec7..0000000 --- a/tests/utils/repository/test_storage_repository.py +++ /dev/null @@ -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') diff --git a/tests/utils/repository/test_training_repository.py b/tests/utils/repository/test_training_repository.py deleted file mode 100644 index e500934..0000000 --- a/tests/utils/repository/test_training_repository.py +++ /dev/null @@ -1,1329 +0,0 @@ -"""Unit tests for TrainingRepository.""" - -from io import BytesIO -from unittest.mock import MagicMock, patch - -import numpy as np -import pandas as pd -import pytest - -from model_manager.sientia.models import LinearRegressionModel -from model_manager.utils.models.train_model_params import TrainModelParams -from model_manager.utils.models.train_model_result import TrainModelResult -from model_manager.utils.repository.training_repository import ( - TrainingRepository, - _apply_support_filters, - _ensure_date_column_parsed, - _frontend_format_to_strftime, -) - - -@pytest.fixture -def mock_logger(): - """Create a mock logger for testing.""" - return MagicMock() - - -@pytest.fixture -def training_repo(mock_logger): - """Create TrainingRepository instance with mock logger.""" - return TrainingRepository(logger=mock_logger) - - -@pytest.fixture -def sample_params(): - """Create sample TrainModelParams for testing.""" - return TrainModelParams( - experiment_run_id=1, - experiment_name='test_experiment', - target_variable='target', - variable_columns=['var1', 'var2', 'var3'], - lag_train={'var1': 0, 'var2': 0, 'var3': 0}, - lag_val={'var1': 0, 'var2': 0, 'var3': 0}, - rem_static_win=False, - low_lim={}, - upp_lim={}, - window=0, - use_scaler=False, - include_ar=False, - train_size=80, - shuffle=True, - bucket_name='test-bucket', - file_name='test.csv', - line_separator=',', - decimal_separator='.', - removed_intervals=[], - model_name='Linear Regression', - degree=1, - interaction_only=False, - nan_treatment='drop', - start_date=None, - end_date=None, - scaler_name='None', - support_filters={}, - static_threshold=None, - date_column=None, - date_format=None, - ) - - -@pytest.fixture -def sample_linear_model(): - """Create a mock LinearRegressionModel with known coefficients.""" - mock_model = MagicMock(spec=LinearRegressionModel) - mock_sklearn_model = MagicMock() - mock_sklearn_model.coef_ = np.array([1.5, -0.75, 2.25]) - mock_sklearn_model.intercept_ = 10.5 - mock_model.regr = mock_sklearn_model - return mock_model - - -class TestExtractModelEquation: - """Tests for _extract_model_equation method.""" - - def test_extract_equation_basic(self, training_repo, sample_params, sample_linear_model): - """Test basic equation extraction with simple coefficients.""" - result = training_repo._extract_model_equation(sample_linear_model, sample_params) - - assert result['target_variable'] == 'target' - assert result['model_type'] == 'Linear Regression' - assert result['intercept'] == 10.5 - - assert 'var1' in result['coefficients'] - assert 'var2' in result['coefficients'] - assert 'var3' in result['coefficients'] - - assert result['coefficients']['var1'] == 1.5 - assert result['coefficients']['var2'] == -0.75 - assert result['coefficients']['var3'] == 2.25 - - def test_extract_equation_string_format( - self, training_repo, sample_params, sample_linear_model - ): - """Test equation string is formatted correctly.""" - result = training_repo._extract_model_equation(sample_linear_model, sample_params) - - expected_string = ( - 'target = 10.500000 + 1.500000 * var1 + -0.750000 * var2 + 2.250000 * var3' - ) - assert result['equation_string'] == expected_string - - def test_extract_equation_latex_format(self, training_repo, sample_params, sample_linear_model): - """Test LaTeX equation is formatted correctly.""" - result = training_repo._extract_model_equation(sample_linear_model, sample_params) - - expected_latex = 'target = 10.500000 + 1.500000 \\cdot var1 + -0.750000 \\cdot var2 + 2.250000 \\cdot var3' - assert result['latex_equation'] == expected_latex - - def test_extract_equation_single_variable(self, training_repo, sample_linear_model): - """Test equation extraction with single variable.""" - params = TrainModelParams( - experiment_run_id=1, - experiment_name='test', - target_variable='y', - variable_columns=['x'], - lag_train={'x': 0}, - lag_val={'x': 0}, - rem_static_win=False, - low_lim={}, - upp_lim={}, - window=0, - use_scaler=False, - include_ar=False, - train_size=80, - shuffle=True, - bucket_name='bucket', - file_name='file.csv', - line_separator=',', - decimal_separator='.', - removed_intervals=[], - model_name='Linear Regression', - degree=1, - interaction_only=False, - nan_treatment='drop', - start_date=None, - end_date=None, - scaler_name='None', - support_filters={}, - static_threshold=None, - date_column=None, - date_format=None, - ) - - # Mock model with single coefficient - mock_model = MagicMock(spec=LinearRegressionModel) - mock_sklearn_model = MagicMock() - mock_sklearn_model.coef_ = np.array([3.14]) - mock_sklearn_model.intercept_ = 2.71 - mock_model.regr = mock_sklearn_model - - result = training_repo._extract_model_equation(mock_model, params) - - assert len(result['coefficients']) == 1 - assert result['coefficients']['x'] == 3.14 - assert result['intercept'] == 2.71 - assert 'y = 2.710000 + 3.140000 * x' == result['equation_string'] - - def test_extract_equation_zero_coefficients(self, training_repo, sample_params): - """Test equation extraction when coefficients are zero.""" - mock_model = MagicMock(spec=LinearRegressionModel) - mock_sklearn_model = MagicMock() - mock_sklearn_model.coef_ = np.array([0.0, 0.0, 0.0]) - mock_sklearn_model.intercept_ = 5.0 - mock_model.regr = mock_sklearn_model - - result = training_repo._extract_model_equation(mock_model, sample_params) - - assert all(v == 0.0 for v in result['coefficients'].values()) - assert result['intercept'] == 5.0 - - def test_extract_equation_negative_intercept(self, training_repo, sample_params): - """Test equation extraction with negative intercept.""" - mock_model = MagicMock(spec=LinearRegressionModel) - mock_sklearn_model = MagicMock() - mock_sklearn_model.coef_ = np.array([1.0, 2.0, 3.0]) - mock_sklearn_model.intercept_ = -5.5 - mock_model.regr = mock_sklearn_model - - result = training_repo._extract_model_equation(mock_model, sample_params) - - assert result['intercept'] == -5.5 - assert 'target = -5.500000 +' in result['equation_string'] - - -class TestInitDataPreprocessor: - """Tests for _init_data_preprocessor method.""" - - def test_init_preprocessor_basic(self, training_repo, sample_params): - """Test basic preprocessor initialization.""" - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.target_variable == 'target' - assert preprocessor.input_columns == ['var1', 'var2', 'var3'] - - def test_init_preprocessor_with_scaler(self, training_repo, sample_params): - """Test preprocessor initialization with scaler enabled.""" - sample_params.use_scaler = True - sample_params.scaler_name = 'Standard Scaler' - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.scaler_name == 'Standard Scaler' - - def test_init_preprocessor_without_scaler(self, training_repo, sample_params): - """Test preprocessor initialization without scaler.""" - sample_params.use_scaler = False - sample_params.scaler_name = 'None' - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.scaler_name == 'None' - - def test_init_preprocessor_with_ar(self, training_repo, sample_params): - """Test preprocessor initialization with autoregressive variable.""" - sample_params.include_ar = True - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.ar_var == 'target' - - def test_init_preprocessor_without_ar(self, training_repo, sample_params): - """Test preprocessor initialization without autoregressive variable.""" - sample_params.include_ar = False - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.ar_var is None - - def test_init_preprocessor_with_static_removal(self, training_repo, sample_params): - """Test preprocessor with static window removal enabled and no static_threshold.""" - sample_params.rem_static_win = True - sample_params.static_threshold = None - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.static_threshold == 1 - - def test_init_preprocessor_with_static_removal_custom_threshold( - self, training_repo, sample_params - ): - """Test preprocessor with static window removal and custom static_threshold.""" - sample_params.rem_static_win = True - sample_params.static_threshold = 500 - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.static_threshold == 500 - - def test_init_preprocessor_without_static_removal_ignores_threshold( - self, training_repo, sample_params - ): - """Test preprocessor without static removal ignores static_threshold.""" - sample_params.rem_static_win = False - sample_params.static_threshold = 500 - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor.static_threshold is None - - def test_init_preprocessor_lag_configuration(self, training_repo, sample_params): - """Test preprocessor lag configuration.""" - sample_params.lag_train = {'var1': 5, 'var2': 5, 'var3': 5} - sample_params.lag_val = {'var1': 3, 'var2': 3, 'var3': 3} - preprocessor = training_repo._init_data_preprocessor(sample_params) - - # Check that lag dictionaries are passed correctly - assert preprocessor.lag_train == {'var1': 5, 'var2': 5, 'var3': 5} - assert preprocessor.lag_transform == {'var1': 3, 'var2': 3, 'var3': 3} - - -class TestGetStaticThreshold: - """Tests for _get_static_threshold method.""" - - def test_get_static_threshold_rem_static_win_false(self, training_repo, sample_params): - """Test returns None when rem_static_win is False.""" - sample_params.rem_static_win = False - sample_params.static_threshold = 500 - - result = training_repo._get_static_threshold(sample_params) - - assert result is None - - def test_get_static_threshold_rem_static_win_true_with_value( - self, training_repo, sample_params - ): - """Test returns static_threshold value when rem_static_win is True and value is set.""" - sample_params.rem_static_win = True - sample_params.static_threshold = 500 - - result = training_repo._get_static_threshold(sample_params) - - assert result == 500 - - def test_get_static_threshold_rem_static_win_true_with_none(self, training_repo, sample_params): - """Test returns 1 when rem_static_win is True and static_threshold is None.""" - sample_params.rem_static_win = True - sample_params.static_threshold = None - - result = training_repo._get_static_threshold(sample_params) - - assert result == 1 - - -class TestInitScalerDict: - """Tests for _init_scaler_dict method.""" - - def test_init_scaler_dict_without_scaler(self, training_repo, sample_params): - """Test scaler dict initialization when scaler is not used.""" - from model_manager.sientia.models import DataPreprocessor - - sample_params.use_scaler = False - process_data = MagicMock(spec=DataPreprocessor) - - result = training_repo._init_scaler_dict(process_data, sample_params) - - assert result == {} - - def test_init_scaler_dict_with_minmax_scaler(self, training_repo, sample_params): - """Test scaler dict initialization with MinMaxScaler.""" - from sientia_do.operations.normalization import MinMaxScaler - - from model_manager.sientia.models import DataPreprocessor - - sample_params.use_scaler = True - - # Create mock MinMaxScaler - mock_scaler = MagicMock(spec=MinMaxScaler) - mock_scaler.x_min = np.array([0.0, 1.0, 2.0]) - mock_scaler.x_max = np.array([10.0, 11.0, 12.0]) - mock_scaler.y_min = 0.5 - mock_scaler.y_max = 100.5 - - # Create mock preprocessor that returns the scaler - process_data = MagicMock(spec=DataPreprocessor) - process_data.get_scaler.return_value = mock_scaler - - result = training_repo._init_scaler_dict(process_data, sample_params) - - # Check feature scalers - assert 'var1' in result - assert 'var2' in result - assert 'var3' in result - assert result['var1'] == {'min': 0.0, 'max': 10.0} - assert result['var2'] == {'min': 1.0, 'max': 11.0} - assert result['var3'] == {'min': 2.0, 'max': 12.0} - - # Check target scaler - assert 'target' in result - assert result['target'] == {'min': 0.5, 'max': 100.5} - - def test_init_scaler_dict_with_z_scaler(self, training_repo, sample_params): - """Test scaler dict initialization with Z_Scaler.""" - from sientia_do.operations.normalization import Z_Scaler - - from model_manager.sientia.models import DataPreprocessor - - sample_params.use_scaler = True - - # Create mock Z_Scaler - mock_scaler = MagicMock(spec=Z_Scaler) - expected_dict = { - 'var1': {'mean': 5.0, 'std': 1.5}, - 'var2': {'mean': 10.0, 'std': 2.0}, - 'target': {'mean': 50.0, 'std': 10.0}, - } - mock_scaler.create_dict.return_value = expected_dict - - # Create mock preprocessor - process_data = MagicMock(spec=DataPreprocessor) - process_data.get_scaler.return_value = mock_scaler - - result = training_repo._init_scaler_dict(process_data, sample_params) - - assert result == expected_dict - mock_scaler.create_dict.assert_called_once() - - -class TestAfterTrainCalculation: - """Tests for after_train_calculation method.""" - - @pytest.fixture - def mock_train_result(self, sample_params, sample_linear_model): - """Create a mock TrainModelResult.""" - x_train = pd.DataFrame( - {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2] - ) - x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) - y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') - y_test = pd.Series([400, 500], index=[3, 4], name='target') - - # Mock predict to return arrays with correct length based on input - def mock_predict(data): - if len(data) == 3: # x_train - return np.array([150.0, 250.0, 350.0]) - else: # x_test - return np.array([450.0, 550.0]) - - sample_linear_model.predict = MagicMock(side_effect=mock_predict) - - return TrainModelResult( - params=sample_params, - process_data=MagicMock(), - x_train=x_train, - x_test=x_test, - y_train=y_train, - y_test=y_test, - regr=sample_linear_model, - scaler_dict={}, - ) - - def test_after_train_adds_predictions(self, training_repo, sample_params, mock_train_result): - """Test that predictions are added to result.""" - result = training_repo.after_train_calculation(sample_params, mock_train_result) - - assert result.y_pred is not None - assert len(result.y_pred) == len(result.y_test) - assert result.y_pred.name == 'target_pred' - - def test_after_train_calculates_metrics(self, training_repo, sample_params, mock_train_result): - """Test that metrics are calculated.""" - result = training_repo.after_train_calculation(sample_params, mock_train_result) - - assert result.mse_val is not None - assert result.mae_val is not None - assert result.r2_val is not None - assert isinstance(result.mse_val, (int, float)) - assert isinstance(result.mae_val, (int, float)) - assert isinstance(result.r2_val, (int, float)) - - def test_after_train_extracts_equation(self, training_repo, sample_params, mock_train_result): - """Test that equation is extracted after training.""" - result = training_repo.after_train_calculation(sample_params, mock_train_result) - - assert result.equation is not None - assert 'target_variable' in result.equation - assert 'coefficients' in result.equation - assert 'intercept' in result.equation - assert 'equation_string' in result.equation - assert 'latex_equation' in result.equation - assert 'model_type' in result.equation - - def test_after_train_sorts_data(self, training_repo, sample_params, mock_train_result): - """Test that data is sorted by index.""" - # Shuffle indices - mock_train_result.x_train = mock_train_result.x_train.sample(frac=1) - mock_train_result.y_train = mock_train_result.y_train.sample(frac=1) - - result = training_repo.after_train_calculation(sample_params, mock_train_result) - - assert result.x_train.index.is_monotonic_increasing - assert result.y_train.index.is_monotonic_increasing - assert result.x_test.index.is_monotonic_increasing - assert result.y_test.index.is_monotonic_increasing - - def test_after_train_logs_success( - self, training_repo, mock_logger, sample_params, mock_train_result - ): - """Test that success is logged.""" - training_repo.after_train_calculation(sample_params, mock_train_result) - - mock_logger.info.assert_called() - assert any( - 'Model metrics calculated successfully' in str(call) - for call in mock_logger.info.call_args_list - ) - - def test_after_train_with_custom_scaler_denormalization( - self, training_repo, sample_params, sample_linear_model - ): - """Test denormalization with custom scaler that has denormalize methods.""" - sample_params.use_scaler = True - - # Create mock data - x_train = pd.DataFrame( - {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2] - ) - x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) - y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') - y_test = pd.Series([400, 500], index=[3, 4], name='target') - - # Mock predict to return arrays with correct length based on input - def mock_predict(data): - if len(data) == 3: # x_train - return np.array([150.0, 250.0, 350.0]) - else: # x_test - return np.array([450.0, 550.0]) - - sample_linear_model.predict = MagicMock(side_effect=mock_predict) - - # Create mock scaler with denormalize methods - mock_scaler = MagicMock() - mock_scaler.denormalize_single_input = MagicMock(side_effect=lambda x, col: x * 2) - - # denormalize_predictions needs to return correct length based on input - def mock_denormalize_predictions(arr, col): - return arr * 2 - - mock_scaler.denormalize_predictions = MagicMock(side_effect=mock_denormalize_predictions) - - # Create mock preprocessor - mock_process_data = MagicMock() - mock_process_data.get_scaler.return_value = mock_scaler - - train_result = TrainModelResult( - params=sample_params, - process_data=mock_process_data, - x_train=x_train, - x_test=x_test, - y_train=y_train, - y_test=y_test, - regr=sample_linear_model, - scaler_dict={}, - ) - - result = training_repo.after_train_calculation(sample_params, train_result) - - # Verify denormalize methods were called - assert mock_scaler.denormalize_single_input.called - assert mock_scaler.denormalize_predictions.called - assert result.y_pred is not None - - def test_after_train_with_sklearn_scaler( - self, training_repo, sample_params, sample_linear_model - ): - """Test denormalization with sklearn StandardScaler.""" - sample_params.use_scaler = True - - # Create mock data - x_train = pd.DataFrame( - {'var1': [1, 2, 3], 'var2': [4, 5, 6], 'var3': [7, 8, 9]}, index=[0, 1, 2] - ) - x_test = pd.DataFrame({'var1': [10, 11], 'var2': [11, 12], 'var3': [12, 13]}, index=[3, 4]) - y_train = pd.Series([100, 200, 300], index=[0, 1, 2], name='target') - y_test = pd.Series([400, 500], index=[3, 4], name='target') - - # Mock predict to return arrays with correct length based on input - def mock_predict(data): - if len(data) == 3: # x_train - return np.array([150.0, 250.0, 350.0]) - else: # x_test - return np.array([450.0, 550.0]) - - sample_linear_model.predict = MagicMock(side_effect=mock_predict) - - # Create mock sklearn scaler (without denormalize methods) - mock_scaler = MagicMock() - # Remove denormalize methods to trigger sklearn path - if hasattr(mock_scaler, 'denormalize_single_input'): - delattr(mock_scaler, 'denormalize_single_input') - mock_scaler.inverse_transform = MagicMock(side_effect=lambda x: x * 2) - - # Create mock preprocessor with feature_names_order - mock_process_data = MagicMock() - mock_process_data.get_scaler.return_value = mock_scaler - mock_process_data.feature_names_order = ['var1', 'var2', 'var3'] - - train_result = TrainModelResult( - params=sample_params, - process_data=mock_process_data, - x_train=x_train, - x_test=x_test, - y_train=y_train, - y_test=y_test, - regr=sample_linear_model, - scaler_dict={}, - ) - - result = training_repo.after_train_calculation(sample_params, train_result) - - # Verify inverse_transform was called - assert mock_scaler.inverse_transform.called - assert result.y_pred is not None - - -class TestTrain: - """Tests for train method.""" - - @pytest.fixture - def sample_csv_data(self): - """Create sample CSV data in BytesIO.""" - csv_content = """var1,var2,var3,target -1.0,2.0,3.0,10.0 -2.0,3.0,4.0,15.0 -3.0,4.0,5.0,20.0 -4.0,5.0,6.0,25.0 -5.0,6.0,7.0,30.0 -6.0,7.0,8.0,35.0 -7.0,8.0,9.0,40.0 -8.0,9.0,10.0,45.0 -9.0,10.0,11.0,50.0 -10.0,11.0,12.0,55.0 -""" - return BytesIO(csv_content.encode('utf-8')) - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_basic_workflow( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test basic training workflow.""" - # Mock load_data to return a DataFrame - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - # Mock split_train_test to return train/test splits - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - result = training_repo.train(sample_csv_data, sample_params) - - # Verify result structure - assert result is not None - assert isinstance(result, TrainModelResult) - assert result.regr is not None - assert result.x_train is not None - assert result.x_test is not None - assert result.y_train is not None - assert result.y_test is not None - assert result.process_data is not None - assert result.scaler_dict is not None - - # Verify load_data was called correctly - mock_load_data.assert_called_once_with( - sample_csv_data, sample_params.line_separator, sample_params.decimal_separator - ) - - # Verify split was called - assert mock_split_train_test.called - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_with_scaler( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test training with scaler enabled.""" - sample_params.use_scaler = True - - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - result = training_repo.train(sample_csv_data, sample_params) - - assert result is not None - assert result.scaler_dict is not None - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_with_shuffle_enabled( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test training with shuffle enabled.""" - sample_params.shuffle = True - - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - training_repo.train(sample_csv_data, sample_params) - - # Verify split was called with shuffle=True - call_kwargs = mock_split_train_test.call_args[1] - assert call_kwargs['shuffle'] is True - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_with_different_train_size( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test training with different train size.""" - sample_params.train_size = 70 - - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - training_repo.train(sample_csv_data, sample_params) - - # Verify split was called with train_size=0.7 - call_kwargs = mock_split_train_test.call_args[1] - assert call_kwargs['train_size'] == 0.7 - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_raises_on_empty_data_after_transform( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test that ValueError is raised when transformed data is empty.""" - # Mock load_data to return empty DataFrame - mock_df = pd.DataFrame( - { - 'var1': [], - 'var2': [], - 'var3': [], - 'target': [], - } - ) - mock_load_data.return_value = mock_df - - with pytest.raises(ValueError, match='Data view is empty after transformation'): - training_repo.train(sample_csv_data, sample_params) - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_logs_success( - self, - mock_load_data, - mock_split_train_test, - training_repo, - mock_logger, - sample_params, - sample_csv_data, - ): - """Test that training success is logged.""" - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - training_repo.train(sample_csv_data, sample_params) - - # Verify success was logged - mock_logger.info.assert_called() - assert any( - 'Model trained successfully' in str(call) for call in mock_logger.info.call_args_list - ) - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_with_custom_separators( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test training with custom line and decimal separators.""" - sample_params.line_separator = ';' - sample_params.decimal_separator = ',' - - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - training_repo.train(sample_csv_data, sample_params) - - # Verify load_data was called with custom separators - mock_load_data.assert_called_once_with(sample_csv_data, ';', ',') - - @patch.object(TrainingRepository, '_configure_datetime_index', lambda self, df, params: df) - @patch('model_manager.utils.repository.training_repository.split_train_test') - @patch('model_manager.utils.repository.training_repository.load_data') - def test_train_result_contains_all_fields( - self, mock_load_data, mock_split_train_test, training_repo, sample_params, sample_csv_data - ): - """Test that TrainModelResult contains all expected fields.""" - mock_df = pd.DataFrame( - { - 'var1': [1, 2, 3, 4, 5], - 'var2': [2, 3, 4, 5, 6], - 'var3': [3, 4, 5, 6, 7], - 'target': [10, 15, 20, 25, 30], - } - ) - mock_load_data.return_value = mock_df - - x_train = pd.DataFrame({'var1': [1, 2, 3], 'var2': [2, 3, 4], 'var3': [3, 4, 5]}) - x_test = pd.DataFrame({'var1': [4, 5], 'var2': [5, 6], 'var3': [6, 7]}) - y_train = pd.Series([10, 15, 20], name='target') - y_test = pd.Series([25, 30], name='target') - mock_split_train_test.return_value = (x_train, x_test, y_train, y_test) - - result = training_repo.train(sample_csv_data, sample_params) - - # Verify all expected fields are present - assert hasattr(result, 'params') - assert hasattr(result, 'process_data') - assert hasattr(result, 'x_train') - assert hasattr(result, 'x_test') - assert hasattr(result, 'y_train') - assert hasattr(result, 'y_test') - assert hasattr(result, 'regr') - assert hasattr(result, 'scaler_dict') - assert result.params == sample_params - - -# ============================================================================ -# Tests for _configure_datetime_index -# ============================================================================ - - -class TestConfigureDatetimeIndex: - """Tests for _configure_datetime_index method.""" - - @pytest.fixture - def training_repo(self, mock_logger): - """Create a TrainingRepository instance.""" - return TrainingRepository(logger=mock_logger) - - @pytest.fixture - def datetime_params(self): - """Minimal params for _configure_datetime_index (date_column/date_format can be None).""" - return TrainModelParams( - experiment_run_id=1, - experiment_name='test', - target_variable='y', - variable_columns=['x'], - lag_train={'x': 0}, - lag_val={'x': 0}, - rem_static_win=False, - low_lim={}, - upp_lim={}, - window=0, - use_scaler=False, - include_ar=False, - train_size=80, - shuffle=True, - bucket_name='b', - file_name='f.csv', - line_separator=',', - decimal_separator='.', - removed_intervals=[], - model_name='Linear Regression', - degree=1, - interaction_only=False, - nan_treatment='drop', - start_date=None, - end_date=None, - scaler_name='None', - support_filters={}, - static_threshold=None, - date_column=None, - date_format=None, - ) - - def test_configure_datetime_index_already_datetime(self, training_repo, datetime_params): - """Test _configure_datetime_index when index is already DatetimeIndex.""" - data = pd.DataFrame( - {'var1': [1, 2, 3], 'var2': [4, 5, 6]}, - index=pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-03']), - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - assert isinstance(result.index, pd.DatetimeIndex) - assert len(result) == 3 - - def test_configure_datetime_index_with_timestamp_column(self, training_repo, datetime_params): - """Test _configure_datetime_index with 'timestamp' column.""" - data = pd.DataFrame( - { - 'timestamp': ['2023-01-01', '2023-01-02', '2023-01-03'], - 'var1': [1, 2, 3], - 'var2': [4, 5, 6], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - assert isinstance(result.index, pd.DatetimeIndex) - assert 'timestamp' not in result.columns - - def test_configure_datetime_index_with_date_column(self, training_repo, datetime_params): - """Test _configure_datetime_index with 'date' column.""" - data = pd.DataFrame( - { - 'date': ['2023-01-01', '2023-01-02', '2023-01-03'], - 'var1': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - assert isinstance(result.index, pd.DatetimeIndex) - assert 'date' not in result.columns - - def test_configure_datetime_index_with_datetime_column(self, training_repo, datetime_params): - """Test _configure_datetime_index with 'datetime' column.""" - data = pd.DataFrame( - { - 'datetime': ['2023-01-01', '2023-01-02', '2023-01-03'], - 'var1': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - assert isinstance(result.index, pd.DatetimeIndex) - assert 'datetime' not in result.columns - - def test_configure_datetime_index_first_column_datetime(self, training_repo, datetime_params): - """Test _configure_datetime_index when first column looks like datetime.""" - data = pd.DataFrame( - { - 'my_date': ['2023-01-01', '2023-01-02', '2023-01-03'], - 'var1': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - assert isinstance(result.index, pd.DatetimeIndex) - assert 'my_date' not in result.columns - - @pytest.mark.filterwarnings('ignore::UserWarning') - def test_configure_datetime_index_no_timestamp_column(self, training_repo, datetime_params): - """Test _configure_datetime_index when no timestamp column found.""" - data = pd.DataFrame( - { - 'var1': ['text_a', 'text_b', 'text_c'], - 'var2': ['text_d', 'text_e', 'text_f'], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - # Should return original data unchanged (no valid datetime columns) - assert 'var1' in result.columns - assert 'var2' in result.columns - - @pytest.mark.filterwarnings('ignore::UserWarning') - def test_configure_datetime_index_invalid_timestamp_column( - self, training_repo, datetime_params - ): - """Test _configure_datetime_index with invalid timestamp values.""" - data = pd.DataFrame( - { - 'timestamp': ['not_a_date', 'also_not', 'nope'], - 'var1': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - # Should skip invalid column and try first column - assert 'var1' in result.columns - - @pytest.mark.filterwarnings('ignore::UserWarning') - def test_configure_datetime_index_invalid_first_column(self, training_repo, datetime_params): - """Test _configure_datetime_index when first column is not datetime.""" - data = pd.DataFrame( - { - 'var1': ['a', 'b', 'c'], - 'var2': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - # Should return original data unchanged - assert 'var1' in result.columns - assert 'var2' in result.columns - - def test_configure_datetime_index_first_column_all_nan(self, training_repo, datetime_params): - """Test _configure_datetime_index when first column has all NaN values.""" - data = pd.DataFrame( - { - 'first_col': [np.nan, np.nan, np.nan], - 'var1': [1, 2, 3], - } - ) - - result = training_repo._configure_datetime_index(data, datetime_params) - - # Should return original data unchanged (first column has no valid values) - assert 'first_col' in result.columns - assert 'var1' in result.columns - - -# ============================================================================ -# Tests for _init_data_preprocessor with removed_intervals -# ============================================================================ - - -class TestInitDataPreprocessorWithRemovedIntervals: - """Tests for _init_data_preprocessor with removed_intervals.""" - - @pytest.fixture - def training_repo(self, mock_logger): - """Create a TrainingRepository instance.""" - return TrainingRepository(logger=mock_logger) - - def test_init_data_preprocessor_with_removed_intervals(self, training_repo, sample_params): - """Test _init_data_preprocessor with removed_intervals.""" - sample_params.removed_intervals = [ - ['2023-01-01', '2023-01-02'], - ['2023-02-01', '2023-02-02'], - ] - - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor is not None - assert preprocessor.removed_intervals is not None - assert len(preprocessor.removed_intervals) == 2 - - def test_init_data_preprocessor_with_tuple_intervals(self, training_repo, sample_params): - """Test _init_data_preprocessor with tuple intervals.""" - sample_params.removed_intervals = [('2023-01-01', '2023-01-02')] - - preprocessor = training_repo._init_data_preprocessor(sample_params) - - assert preprocessor is not None - assert preprocessor.removed_intervals is not None - - -# ============================================================================ -# Tests for _extract_model_equation with polynomial features -# ============================================================================ - - -class TestExtractModelEquationPolynomial: - """Tests for _extract_model_equation with polynomial features.""" - - @pytest.fixture - def training_repo(self, mock_logger): - """Create a TrainingRepository instance.""" - return TrainingRepository(logger=mock_logger) - - def test_extract_model_equation_polynomial(self, training_repo, sample_params): - """Test _extract_model_equation with polynomial features.""" - # Create a mock regressor with polynomial features - mock_regr = MagicMock() - mock_regr.regr.coef_ = np.array([0.5, 0.3, 0.2]) - mock_regr.regr.intercept_ = 1.0 - mock_regr.poly_feature_names = ['var1', 'var2', 'var1^2'] - - sample_params.degree = 2 - - result = training_repo._extract_model_equation(mock_regr, sample_params) - - assert 'equation_string' in result - assert 'latex_equation' in result - assert 'var1' in result['equation_string'] - - def test_extract_model_equation_linear(self, training_repo, sample_params): - """Test _extract_model_equation with linear features.""" - mock_regr = MagicMock() - mock_regr.regr.coef_ = np.array([0.5, 0.3]) - mock_regr.regr.intercept_ = 1.0 - mock_regr.poly_feature_names = None - - sample_params.degree = 1 - - result = training_repo._extract_model_equation(mock_regr, sample_params) - - assert 'equation_string' in result - assert 'latex_equation' in result - - -# ============================================================================ -# Tests for _frontend_format_to_strftime and _ensure_date_column_parsed -# ============================================================================ - - -class TestFrontendFormatToStrftime: - """Tests for _frontend_format_to_strftime (alias from models).""" - - def test_empty_or_none_returns_none(self): - """Empty string or None returns None (same behavior as models module).""" - assert _frontend_format_to_strftime('') is None - assert _frontend_format_to_strftime(None) is None - - def test_allowed_formats_convert_correctly(self): - """Allowed frontend formats convert to strftime via mapper.""" - assert _frontend_format_to_strftime('dd/MM/yyyy HH:mm:ss') == '%d/%m/%Y %H:%M:%S' - assert _frontend_format_to_strftime('yyyy-MM-dd HH:mm:ss') == '%Y-%m-%d %H:%M:%S' - assert _frontend_format_to_strftime('MM/dd/yyyy HH:mm:ss') == '%m/%d/%Y %H:%M:%S' - - def test_invalid_format_raises(self): - """Invalid date format raises ValueError.""" - with pytest.raises(ValueError, match='Invalid date_format'): - _frontend_format_to_strftime('yyyy-MM-dd') - - -class TestEnsureDateColumnParsed: - """Tests for _ensure_date_column_parsed.""" - - @pytest.fixture - def date_params(self): - """Params with date_column and date_format set.""" - return TrainModelParams( - experiment_run_id=1, - experiment_name='test', - target_variable='y', - variable_columns=['x'], - lag_train={'x': 0}, - lag_val={'x': 0}, - rem_static_win=False, - low_lim={}, - upp_lim={}, - window=0, - use_scaler=False, - include_ar=False, - train_size=80, - shuffle=True, - bucket_name='b', - file_name='f.csv', - line_separator=',', - decimal_separator='.', - removed_intervals=[], - model_name='Linear Regression', - degree=1, - interaction_only=False, - nan_treatment='drop', - start_date=None, - end_date=None, - scaler_name='None', - support_filters={}, - static_threshold=None, - date_column='ts', - date_format='yyyy-MM-dd HH:mm:ss', - ) - - def test_returns_unchanged_when_no_date_column(self, date_params): - """When params.date_column is None, data is returned unchanged.""" - date_params.date_column = None - date_params.date_format = None - data = pd.DataFrame({'ts': ['2023-01-01'], 'x': [1]}) - result = _ensure_date_column_parsed(data, date_params) - pd.testing.assert_frame_equal(result, data) - - def test_returns_unchanged_when_column_missing(self, date_params): - """When date_column not in data columns, data is returned unchanged.""" - data = pd.DataFrame({'other': [1], 'x': [2]}) - result = _ensure_date_column_parsed(data, date_params) - pd.testing.assert_frame_equal(result, data) - - def test_parses_column_with_format(self, date_params): - """When date_column and date_format set, column is parsed as datetime.""" - data = pd.DataFrame( - { - 'ts': ['2023-01-01 10:00:00', '2023-06-15 14:30:00'], - 'x': [1, 2], - } - ) - result = _ensure_date_column_parsed(data, date_params) - assert result['ts'].dtype == 'datetime64[ns]' - assert result['ts'].iloc[0].year == 2023 - assert result['ts'].iloc[0].month == 1 - assert result['ts'].iloc[1].month == 6 - - def test_invalid_values_coerced_to_nat(self, date_params): - """Invalid date strings are coerced to NaT when format is set.""" - date_params.date_format = 'yyyy-MM-dd HH:mm:ss' - data = pd.DataFrame( - { - 'ts': ['2023-01-01 00:00:00', 'not-a-date', '2023-12-31 00:00:00'], - 'x': [1, 2, 3], - } - ) - result = _ensure_date_column_parsed(data, date_params) - assert pd.isna(result['ts'].iloc[1]) - assert result['ts'].iloc[0].year == 2023 - assert result['ts'].iloc[2].month == 12 - - -# ============================================================================ -# Tests for _apply_support_filters -# ============================================================================ - - -class TestApplySupportFilters: - """Tests for _apply_support_filters function.""" - - def test_empty_support_filters_returns_unchanged(self): - """When support_filters is empty, data_view is returned unchanged.""" - data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) - result = _apply_support_filters(data, 'target', {}) - pd.testing.assert_frame_equal(result, data) - - def test_target_not_in_columns_returns_unchanged(self): - """When target_variable is not in data_view columns, return unchanged.""" - data = pd.DataFrame({'x': [1, 2, 3], 'y': [10, 20, 30]}) - result = _apply_support_filters(data, 'target', {'x': {}}) - pd.testing.assert_frame_equal(result, data) - - def test_variable_not_in_columns_skipped(self): - """When a filter variable is not in data_view, that variable is skipped.""" - data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) - support_filters = { - 'missing_var': { - 'upper_line': {'intercept': 100, 'angle': 10}, - 'lower_line': {'intercept': 0, 'angle': -10}, - }, - } - result = _apply_support_filters(data, 'target', support_filters) - pd.testing.assert_frame_equal(result, data) - - def test_snake_case_upper_lower_line(self): - """Support filters with upper_line/lower_line (snake_case) filter rows.""" - data = pd.DataFrame( - { - 'x': [1.0, 2.0, 3.0, 4.0], - 'target': [2.0, 4.0, 6.0, 8.0], - } - ) - support_filters = { - 'x': { - 'upper_line': {'intercept': 1.0, 'angle': 50}, - 'lower_line': {'intercept': -1.0, 'angle': -50}, - }, - } - result = _apply_support_filters(data, 'target', support_filters) - assert len(result) <= 4 - assert list(result.columns) == ['x', 'target'] - - def test_camel_case_upper_lower_line(self): - """Support filters with upperLine/lowerLine (camelCase) are accepted.""" - data = pd.DataFrame( - { - 'x': [1.0, 2.0, 3.0], - 'target': [1.0, 2.0, 3.0], - } - ) - support_filters = { - 'x': { - 'upperLine': {'intercept': 2, 'angle': 5}, - 'lowerLine': {'intercept': 0, 'angle': -5}, - }, - } - result = _apply_support_filters(data, 'target', support_filters) - assert len(result) <= 3 - assert list(result.columns) == ['x', 'target'] - - def test_two_variables_ands_masks(self): - """Two variables apply AND of both masks.""" - data = pd.DataFrame( - { - 'a': [1.0, 2.0, 3.0], - 'b': [1.0, 2.0, 3.0], - 'target': [2.0, 2.0, 2.0], - } - ) - support_filters = { - 'a': { - 'upper_line': {'intercept': 10, 'angle': 45}, - 'lower_line': {'intercept': -10, 'angle': -45}, - }, - 'b': { - 'upper_line': {'intercept': 10, 'angle': 45}, - 'lower_line': {'intercept': -10, 'angle': -45}, - }, - } - result = _apply_support_filters(data, 'target', support_filters) - assert len(result) <= 3 - assert list(result.columns) == ['a', 'b', 'target'] - - def test_missing_upper_or_lower_skips_variable(self): - """If upper_line or lower_line is missing, that variable is skipped.""" - data = pd.DataFrame({'x': [1, 2, 3], 'target': [10, 20, 30]}) - support_filters = { - 'x': {'upper_line': {'intercept': 100, 'angle': 0}}, - } - result = _apply_support_filters(data, 'target', support_filters) - pd.testing.assert_frame_equal(result, data) diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 89e373e..35170b7 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -20,6 +20,10 @@ def mock_env_vars(): 'PROJECT_NAME': 'test-project', 'TRAIN_TASK_QUEUE': 'train_model-local_queue', 'CLEANUP_TASK_QUEUE': 'cleanup-local_queue', + 'RUNTIME': 'model-manager-worker', + 'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local', + 'STORE_OWNER': 'sientia', + 'STORE_REPO': 'model-library-store', } with patch.dict(os.environ, env_vars, clear=False): @@ -176,6 +180,8 @@ def test_start_prometheus_server_failure( @patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') @patch('model_manager.worker.worker.build_mongodb_config') @patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_mlflow_config') @@ -191,6 +197,8 @@ async def test_main_successful_startup( mock_build_mlflow, mock_build_postgres, mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, mock_notification_handler_class, mock_activities_class, mock_runtime_class, @@ -220,6 +228,23 @@ async def test_main_successful_startup( mock_notification_handler_class.return_value = mock_notification_handler mock_activities_class.return_value = mock_activities + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + mock_runtime = Mock() mock_runtime_class.return_value = mock_runtime @@ -262,6 +287,8 @@ async def test_main_successful_startup( @patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') @patch('model_manager.worker.worker.build_mongodb_config') @patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_mlflow_config') @@ -277,6 +304,8 @@ async def test_main_handles_exception( mock_build_mlflow, mock_build_postgres, mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, mock_notification_handler_class, mock_activities_class, mock_runtime_class, @@ -320,6 +349,23 @@ async def test_main_handles_exception( mock_app_up = Mock() mock_metrics.APP_UP.labels.return_value = mock_app_up + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + # Run main() and expect SystemExit with pytest.raises(SystemExit) as exc_info: await main() @@ -343,6 +389,8 @@ async def test_main_handles_exception( @patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') @patch('model_manager.worker.worker.build_mongodb_config') @patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_mlflow_config') @@ -358,6 +406,8 @@ async def test_main_temporal_client_configuration( mock_build_mlflow, mock_build_postgres, mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, mock_notification_handler_class, mock_activities_class, mock_runtime_class, @@ -377,6 +427,10 @@ async def test_main_temporal_client_configuration( 'TEMPORAL_HOST': 'temporal.example.com:7233', 'TEMPORAL_NAMESPACE': 'production', 'TEMPORAL_USE_TLS': 'true', + 'RUNTIME': 'model-manager-worker', + 'STORE_BASE_URL': 'http://sientia-plugin-store.svc.cluster.local', + 'STORE_OWNER': 'sientia', + 'STORE_REPO': 'model-library-store', }, ): # Setup mocks @@ -390,6 +444,23 @@ async def test_main_temporal_client_configuration( mock_build_mlflow.return_value = {} mock_build_minio.return_value = {} + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + mock_notification_handler = Mock() mock_notification_handler.shutdown = Mock() mock_notification_handler_class.return_value = mock_notification_handler @@ -431,6 +502,8 @@ async def test_main_temporal_client_configuration( @patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') @patch('model_manager.worker.worker.build_mongodb_config') @patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_mlflow_config') @@ -446,6 +519,8 @@ async def test_main_worker_configuration( mock_build_mlflow, mock_build_postgres, mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, mock_notification_handler_class, mock_activities_class, mock_runtime_class, @@ -500,6 +575,23 @@ async def test_main_worker_configuration( mock_app_up = Mock() mock_metrics.APP_UP.labels.return_value = mock_app_up + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + # Run main() with pytest.raises(SystemExit): await main() @@ -539,6 +631,8 @@ async def test_main_worker_configuration( @patch('model_manager.worker.worker.Runtime') @patch('model_manager.worker.worker.Activities') @patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') @patch('model_manager.worker.worker.build_mongodb_config') @patch('model_manager.worker.worker.build_postgres_config') @patch('model_manager.worker.worker.build_mlflow_config') @@ -554,6 +648,8 @@ async def test_main_schedule_creation_failure_does_not_stop_worker( mock_build_mlflow, mock_build_postgres, mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, mock_notification_handler_class, mock_activities_class, mock_runtime_class, @@ -603,6 +699,106 @@ async def test_main_schedule_creation_failure_does_not_stop_worker( mock_app_up = Mock() mock_metrics.APP_UP.labels.return_value = mock_app_up + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + + +@pytest.mark.asyncio +@patch('model_manager.worker.worker.Worker') +@patch('model_manager.worker.worker.client.Client') +@patch('model_manager.worker.worker.Runtime') +@patch('model_manager.worker.worker.Activities') +@patch('model_manager.worker.worker.NotificationHandler') +@patch('model_manager.worker.worker.PluginStore') +@patch('model_manager.worker.worker.build_plugin_store_config') +@patch('model_manager.worker.worker.build_mongodb_config') +@patch('model_manager.worker.worker.build_postgres_config') +@patch('model_manager.worker.worker.build_mlflow_config') +@patch('model_manager.worker.worker.build_minio_config') +@patch('model_manager.worker.worker.get_logger') +@patch('model_manager.worker.worker.start_prometheus_server') +@patch('model_manager.worker.worker.metrics') +async def test_main_missing_runtime_fails_fast( + mock_metrics, + mock_start_prometheus, + mock_get_logger, + mock_build_minio, + mock_build_mlflow, + mock_build_postgres, + mock_build_mongodb, + mock_build_plugin_store_config, + mock_plugin_store_class, + mock_notification_handler_class, + mock_activities_class, + mock_runtime_class, + mock_client_class, + mock_worker_class, + mock_logger, +): + """Test that main() fails fast when RUNTIME is missing.""" + from model_manager.worker.worker import main + + mock_get_logger.return_value = mock_logger + mock_build_mongodb.return_value = { + 'connection_string': 'mongodb://test', + 'database_name': 'test_db', + 'uri': 'localhost:27018', + } + mock_build_postgres.return_value = {} + mock_build_mlflow.return_value = {} + mock_build_minio.return_value = {} + + mock_notification_handler = Mock() + mock_notification_handler.shutdown = Mock() + mock_notification_handler_class.return_value = mock_notification_handler + + mock_activities = AsyncMock() + mock_activities.shutdown = AsyncMock() + mock_activities_class.return_value = mock_activities + + mock_app_up = Mock() + mock_metrics.APP_UP.labels.return_value = mock_app_up + + mock_plugin_store_instance = AsyncMock() + mock_plugin_store_instance.install_runtime = AsyncMock( + return_value={'runtime': 'model-manager-worker', 'installed': []}, + ) + mock_plugin_store_class.return_value = mock_plugin_store_instance + mock_build_plugin_store_config.return_value = { + 'base_url': 'http://sientia-plugin-store.svc.cluster.local', + 'owner': 'sientia', + 'repo': 'model-library-store', + 'branch': 'main', + 'username': 'gitea-user', + 'password': 'gitea-password', + 'pypi_index_url': 'http://library-distribution-server.library.svc.cluster.local:5000', + 'pypi_username': None, + 'pypi_password': None, + } + + # Ensure RUNTIME is not defined + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(SystemExit) as exc_info: + await main() + + assert exc_info.value.code == 1 + mock_logger.custom_critical.assert_called_once() + mock_app_up.set.assert_called_with(0) + # Run main() - should not fail despite schedule creation error with pytest.raises(SystemExit): await main() diff --git a/values.yaml b/values.yaml index 92e0af4..7c999f8 100644 --- a/values.yaml +++ b/values.yaml @@ -1,26 +1,222 @@ -# Default values for sientia-module. +# +# Default values for sientia-dataops-model-manager using the sientia-module chart (0.6.x). # This is a YAML-formatted file. # Declare variables to be passed into your templates. +# -# This will set the replicaset count more information can be found here: https://kubernetes.io/docs/concepts/workloads/controllers/replicaset/ -replicaCount: 1 +projectName: &projectName "sientia-dataops-model-manager" -# This sets the container image more information can be found here: https://kubernetes.io/docs/concepts/containers/images/ +# ----------------------------------------------------------------------------- +# Image configuration (chart-level) +# ----------------------------------------------------------------------------- +# The sientia-module chart allows overriding the image used by all runtimes. +# Per requirement, we deploy using the sientia-module image v1.0.0. image: - repository: aignosi.azurecr.io/sientia-dataops-model-manager - # This sets the pull policy for images. + repository: aignosi.azurecr.io/sientia-module pullPolicy: IfNotPresent - # Overrides the image tag whose default is the chart appVersion. - tag: "1.1.2" + tag: "1.0.0" + +# ----------------------------------------------------------------------------- +# Global configuration shared by all runtimes +# ----------------------------------------------------------------------------- +global: + # Namespace used by the chart. + namespace: sientia + + # Common labels applied to pods (can be extended per project). + commonLabels: {} + + # Resources inherited by all runtimes unless overridden. + resources: {} + # We usually recommend not to specify default resources and to leave this as a conscious + # choice for the user. This also increases chances charts run on environments with little + # resources, such as Minikube. If you do want to specify resources, uncomment the following + # lines, adjust them as necessary, and remove the curly braces after 'resources:'. + # limits: + # cpu: 100m + # memory: 128Mi + # requests: + # cpu: 100m + # memory: 128Mi + + # Probes inherited by all runtimes unless overridden. + # More information: + # https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/ + livenessProbe: + exec: + command: + - python3 + - -c + - "import requests; requests.get('http://localhost:9090/metrics')" + initialDelaySeconds: 20 + periodSeconds: 30 + + readinessProbe: + exec: + command: + - python3 + - -c + - "import requests; requests.get('http://localhost:9090/metrics')" + initialDelaySeconds: 10 + periodSeconds: 15 + + # Autoscaling configuration inherited by all runtimes unless overridden. + # More information: + # https://kubernetes.io/docs/concepts/workloads/autoscaling/ + autoscaling: + enabled: false + minReplicas: 1 + maxReplicas: 100 + targetCPUUtilizationPercentage: 80 + # targetMemoryUtilizationPercentage: 80 + + # Environment variables shared by all runtimes. + env: + - name: POSTGRES_HOST + value: "paradedb-rw.paradedb.svc.cluster.local" + - name: POSTGRES_PORT + value: "5432" + - name: POSTGRES_USER + value: "postgres" + - name: POSTGRES_PASSWORD + value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3" + - name: POSTGRES_DBNAME + value: "sientia-core-mlops-bff" + - name: POSTGRES_MIN_CONNECTIONS + value: "10" + - name: POSTGRES_MAX_CONNECTIONS + value: "30" + + - name: MLFLOW_URL + value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local:80" + - name: MLFLOW_USERNAME + value: "aignosi" + - name: MLFLOW_PASSWORD + value: "1L0FP50j3ncp123" + + - name: LOG_LEVEL + value: "DEBUG" + - name: HTTP_METRICS_PORT + value: "9090" + - name: HTTP_SDK_METRICS_PORT + value: "9091" + + - name: TEMPORAL_HOST + value: "temporal-frontend.temporal.svc.cluster.local:7233" + - name: TEMPORAL_NAMESPACE + value: "model-manager" + - name: TRAIN_TASK_QUEUE + value: "train_model-queue" + - name: CLEANUP_TASK_QUEUE + value: "cleanup-queue" + - name: TEMPORAL_USE_TLS + value: "false" + + - name: STORE_BASE_URL + value: "http://gitea-http.gitea.svc.cluster.local" + - name: STORE_OWNER + value: "aignosi" + - name: STORE_REPO + value: "suse-model-store" + - name: STORE_USERNAME + valueFrom: + secretKeyRef: + name: sientia-plugin-store-credentials + key: username + - name: STORE_PASSWORD + valueFrom: + secretKeyRef: + name: sientia-plugin-store-credentials + key: password + - name: STORE_CACHE_TTL_SECONDS + value: "3600" + + - name: MONGODB_USERNAME + value: "root" + - name: MONGODB_PASSWORD + value: "wKZDbMNU1c" + - name: MONGODB_URL + value: "my-release-mongodb.mongodb.svc.cluster.local:27017" + - name: MONGODB_DATABASE + value: "sientia" + - name: MONGODB_TTL_INDEX_HOURS + value: "1" + + - name: MINIO_ENDPOINT_URL + value: "http://minio.minio.svc.cluster.local:9000" + - name: MINIO_ACCESS_KEY + value: "model-training-user" + - name: MINIO_SECRET_KEY + value: "modelTrainingUser123" + - name: MINIO_REGION + value: "us-east-1" + - name: MINIO_USE_SSL + value: "false" + - name: MINIO_MAX_RETRY_ATTEMPTS + value: "3" + - name: MINIO_RETRY_MODE + value: "adaptive" + - name: MINIO_CONNECT_TIMEOUT + value: "10" + - name: MINIO_READ_TIMEOUT + value: "60" + + - name: TIMEOUT_VALIDATE_PARAMS + value: "30" + - name: TIMEOUT_TRAIN_MODEL + value: "2700" + - name: TIMEOUT_DELETE_FILE + value: "120" + - name: TIMEOUT_UPDATE_DATABASE + value: "30" + + - name: CLEANUP_RETENTION_HOURS + value: "24" + - name: CLEANUP_DRY_RUN + value: "false" + - name: TIMEOUT_CLEANUP_MINIO + value: "300" + - name: TIMEOUT_CLEANUP_LOCAL + value: "120" + - name: MAX_KEYS_CLEANUP + value: "1000" + - name: DEFAULT_CLEANUP_BUCKET + value: "model-training" + + # Cleanup Schedule Configuration + - name: CLEANUP_SCHEDULE_ID + value: "cleanup-files-daily" + - name: CLEANUP_CRON + value: "0 0 * * *" # Midnight UTC + - name: CLEANUP_TIMEZONE + value: "UTC" + - name: CLEANUP_EXECUTION_TIMEOUT_HOURS + value: "1" + + - name: PYPI_SERVER + value: "http://library-distribution-server.library.svc.cluster.local:5000" + +# ----------------------------------------------------------------------------- +# Runtimes configuration +# ----------------------------------------------------------------------------- +# Each runtime inherits settings from `global` (resources, env, probes, autoscaling) +# unless overridden here. +runtimes: + - name: "model-manager-worker" + # Replicas for this runtime. Replaces the old replicaCount. + replicas: 1 + +# ----------------------------------------------------------------------------- +# Chart-level configuration (applies to all runtimes) +# ----------------------------------------------------------------------------- # This is for the secrets for pulling an image from a private repository more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/pull-image-private-registry/ imagePullSecrets: -- name: docker-hub-secret + - name: docker-hub-secret # This is to override the chart name. -nameOverride: "sientia-dataops-model-manager" -fullnameOverride: "sientia-dataops-model-manager" -namespace: sientia +nameOverride: *projectName +fullnameOverride: *projectName # This section builds out the service account more information can be found here: https://kubernetes.io/docs/concepts/security/service-accounts/ serviceAccount: @@ -32,7 +228,7 @@ serviceAccount: annotations: {} # The name of the service account to use. # If not set and create is true, a name is generated using the fullname template - name: "sientia-dataops-model-manager" + name: *projectName # This is for setting Kubernetes Annotations to a Pod. # For more information checkout: https://kubernetes.io/docs/concepts/overview/working-with-objects/annotations/ @@ -52,45 +248,6 @@ securityContext: {} # runAsNonRoot: true # runAsUser: 1000 -resources: {} - # We usually recommend not to specify default resources and to leave this as a conscious - # choice for the user. This also increases chances charts run on environments with little - # resources, such as Minikube. If you do want to specify resources, uncomment the following - # lines, adjust them as necessary, and remove the curly braces after 'resources:'. - # limits: - # cpu: 100m - # memory: 128Mi - # requests: - # cpu: 100m - # memory: 128Mi - -# This is to setup the liveness and readiness probes more information can be found here: https://kubernetes.io/docs/tasks/configure-pod-container/configure-liveness-readiness-startup-probes/ -livenessProbe: - exec: - command: - - python3 - - -c - - "import requests; requests.get('http://localhost:9090/metrics')" - initialDelaySeconds: 20 - periodSeconds: 30 - -readinessProbe: - exec: - command: - - python3 - - -c - - "import requests; requests.get('http://localhost:9090/metrics')" - initialDelaySeconds: 10 - periodSeconds: 15 - -# This section is for setting up autoscaling more information can be found here: https://kubernetes.io/docs/concepts/workloads/autoscaling/ -autoscaling: - enabled: false - minReplicas: 1 - maxReplicas: 100 - targetCPUUtilizationPercentage: 80 - # targetMemoryUtilizationPercentage: 80 - # Additional volumes on the output Deployment definition. volumes: - name: reports-volume @@ -151,119 +308,6 @@ serviceMonitor: additionalLabels: release: kube-prometheus-stack -env: - - name: POSTGRES_HOST - value: "paradedb-rw.paradedb.svc.cluster.local" - - name: POSTGRES_PORT - value: "5432" - - name: POSTGRES_USER - value: "postgres" - - name: POSTGRES_PASSWORD - value: "nFqc81y6kwmr2zuAIx43DhiOosFCVPpeEfTtTWZflkNjB2j1KtEeIANkhFR9mAX3" - - name: POSTGRES_DBNAME - value: "sientia-core-mlops-bff" - - name: POSTGRES_MIN_CONNECTIONS - value: "10" - - name: POSTGRES_MAX_CONNECTIONS - value: "30" - - - name: MLFLOW_URL - value: "http://sientia-tracker-mlflow-tracking.sientia-tracker.svc.cluster.local:80" - - name: MLFLOW_USERNAME - value: "aignosi" - - name: MLFLOW_PASSWORD - value: "1L0FP50j3ncp123" - - - name: LOG_LEVEL - value: "DEBUG" - - name: HTTP_METRICS_PORT - value: "9090" - - name: HTTP_SDK_METRICS_PORT - value: "9091" - - name: PROJECT_NAME - value: "sientia-model-manager" - - - name: TEMPORAL_HOST - value: "temporal-frontend.temporal.svc.cluster.local:7233" - - name: TEMPORAL_NAMESPACE - value: "model-manager" - - name: TRAIN_TASK_QUEUE - value: "train_model-queue" - - name: CLEANUP_TASK_QUEUE - value: "cleanup-queue" - - name: TEMPORAL_USE_TLS - value: "false" - - - name: MONGODB_USERNAME - value: "root" - - name: MONGODB_PASSWORD - value: "wKZDbMNU1c" - - name: MONGODB_URL - value: "my-release-mongodb.mongodb.svc.cluster.local:27017" - - name: MONGODB_DATABASE - value: "sientia" - - name: MONGODB_TTL_INDEX_HOURS - value: "1" - - - name: MINIO_ENDPOINT_URL - value: "http://minio.minio.svc.cluster.local:9000" - - name: MINIO_ACCESS_KEY - value: "model-training-user" - - name: MINIO_SECRET_KEY - value: "modelTrainingUser123" - - name: MINIO_REGION - value: "us-east-1" - - name: MINIO_USE_SSL - value: "false" - - name: MINIO_MAX_RETRY_ATTEMPTS - value: "3" - - name: MINIO_RETRY_MODE - value: "adaptive" - - name: MINIO_CONNECT_TIMEOUT - value: "10" - - name: MINIO_READ_TIMEOUT - value: "60" - - - name: TIMEOUT_VALIDATE_PARAMS - value: "30" - - name: TIMEOUT_TRAIN_MODEL - value: "2700" - - name: TIMEOUT_DELETE_FILE - value: "120" - - name: TIMEOUT_UPDATE_DATABASE - value: "30" - - - name: CLEANUP_RETENTION_HOURS - value: "24" - - name: CLEANUP_DRY_RUN - value: "false" - - name: TIMEOUT_CLEANUP_MINIO - value: "300" - - name: TIMEOUT_CLEANUP_LOCAL - value: "120" - - name: MAX_KEYS_CLEANUP - value: "1000" - - name: DEFAULT_CLEANUP_BUCKET - value: "model-training" - - # Cleanup Schedule Configuration - - name: CLEANUP_SCHEDULE_ID - value: "cleanup-files-daily" - - name: CLEANUP_CRON - value: "0 0 * * *" # Midnight UTC - - name: CLEANUP_TIMEZONE - value: "UTC" - - name: CLEANUP_EXECUTION_TIMEOUT_HOURS - value: "1" - - - name: EXTRA_PIP_REQUIREMENTS - value: "git+https://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@github.com/Aignosi/sientia-mlops-library.git" - - - name: POD_ID - valueFrom: - fieldRef: - fieldPath: metadata.name - ssh: enabled: false secretName: git-ssh-key-sientia-model-manager-worker @@ -308,11 +352,21 @@ grafanaDatasource: # jsonData: # timeInterval: "5s" -# kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password=5I5zpQ6sRaHqX1hD3dr+2mo647yO3FRc359/wu6gsP+ACRDRz5mp - -# helm upgrade --install sientia-dataops-model-manager sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.0 - +# ----------------------------------------------------------------------------- +# Helm usage examples +# ----------------------------------------------------------------------------- +# kubectl create secret docker-registry docker-hub-secret --namespace sientia --docker-server=http://aignosi.azurecr.io --docker-username=aignosi --docker-password= +# +# helm upgrade --install sientia-dataops-model-manager sientia/sientia-module -n sientia --create-namespace -f ./values.yaml --version 0.6.1 +# +# Global/runtimes layout note: +# - Shared configuration lives under `global` (env, probes, autoscaling, namespace). +# - Individual runtimes are defined under `runtimes`, each with its own `name` and `replicas`. +# - Runtimes inherit `global` settings unless overridden at the runtime level. +# # kubectl create secret generic git-ssh-key-sientia-model-manager-worker \ # --namespace sientia \ # --from-file=ssh-privatekey=git_key \ # --type=kubernetes.io/ssh-auth + +