418 lines
16 KiB
Markdown
418 lines
16 KiB
Markdown
# Train Model Workflow IO Diff (`main` vs current branch)
|
|
|
|
Base comparison: `git diff main...HEAD`
|
|
Workflow analyzed: `train_model`
|
|
|
|
## 1) Executive overview
|
|
|
|
This branch introduces a structural refactor of the training stack and a contract update for workflow input/output.
|
|
|
|
Main impacts:
|
|
|
|
- The old in-house training stack (`TrainingRepository`, `ModelRepository`, `StorageRepository`, `model_manager.sientia.models`) was replaced by:
|
|
- `DataManagerRepository` (data prep + metrics + report generation)
|
|
- `SientiaModel` wrapper from plugin store (`sientia_model`)
|
|
- `SientiaMLflowRepository` (MLflow integration)
|
|
- `MinioRepository` (storage integration)
|
|
- Input contract moved from many fixed legacy ML params to a plugin/wrapper-oriented schema (`model_type`, `*_kwargs`, `model_metadata`, optional `val_file_name`).
|
|
- Workflow return changed from `None` to a serializable result object (`dict[str, Any] | None`) containing training execution metadata.
|
|
- Queue naming and worker bootstrap architecture now depend on runtime (`train_model-<runtime>-queue`).
|
|
|
|
---
|
|
|
|
## 2) Input contract diff (before vs now)
|
|
|
|
### 2.1 Previous contract (`main`)
|
|
|
|
`TrainModelParams` in `main` required a large set of explicit fields for the old preprocessing/model pipeline, focused only in linear regression model:
|
|
|
|
- Core:
|
|
- `experiment_run_id`, `variable_columns`, `target_variable`
|
|
- `bucket_name`, `file_name`, `line_separator`, `decimal_separator`
|
|
- `train_size`, `shuffle`
|
|
- Legacy preprocessing/model fields focused only in linear regression model (required in `from_dict`):
|
|
- `lag_train`, `lag_val`
|
|
- `rem_static_win`, `low_lim`, `upp_lim`, `window`
|
|
- `use_scaler`, `include_ar`, `scaler_name`
|
|
- `removed_intervals`, `start_date`, `end_date`, `nan_treatment`
|
|
- `degree`, `interaction_only`
|
|
- `experiment_name`, `model_name`
|
|
- `support_filters` (optional dict), `static_threshold` (optional int)
|
|
|
|
Validation was strongly tied to this structure (lag ranges, limits consistency, polynomial/scaler constraints, etc.).
|
|
|
|
### 2.2 Current contract (this branch)
|
|
|
|
`TrainModelParams` now supports a plugin-driven schema and wrapper kwargs:
|
|
|
|
- Kept/mandatory core fields:
|
|
- `experiment_run_id` (now accepts numeric string too; coerced to int)
|
|
- `variable_columns`, `target_variable`
|
|
- `bucket_name`, `file_name`, `line_separator`, `decimal_separator`
|
|
- `train_size`, `shuffle`
|
|
- `model_name`
|
|
- `model_type`
|
|
- `data_model_kwargs`, `model_kwargs`, `opt_params` (required as dict by current `from_dict`)
|
|
- New/updated fields:
|
|
- `random_state` (default `42`)
|
|
- `val_file_name` (optional explicit validation file)
|
|
- `model_id` (currently optional, but needs discussion, since the model metadata in MongoDB should be created before the model training)
|
|
- Removed from required input contract:
|
|
- `lag_train`, `lag_val`, `rem_static_win`, `low_lim`, `upp_lim`, `window`
|
|
- `use_scaler`, `include_ar`
|
|
- `degree`, `interaction_only`, `nan_treatment`
|
|
- `start_date`, `end_date`, `scaler_name`
|
|
- `removed_intervals`, `support_filters`, `static_threshold`
|
|
- Parameters internally derived:
|
|
- `model_metadata` model type info from plugin store.
|
|
- `run_name` is internally derived from experiment name and datetime.
|
|
- `experiment_name` is internally derived from `model_name`.
|
|
|
|
### 2.3 Validation behavior changes
|
|
|
|
Before:
|
|
- Validation was mostly hardcoded business checks tied to legacy linear/polynomial stack.
|
|
|
|
Now:
|
|
- Validation still checks core constraints (`train_size`, non-empty strings, etc.), but model-specific validation moved to JSON Schema driven checks, using OpenAPI/JSON Schema definitions from plugin store:
|
|
- `model_metadata.schemas.components.schemas.data_model`
|
|
- `model_metadata.schemas.components.schemas.model`
|
|
- `model_metadata.schemas.components.schemas.opt_params`
|
|
- `model_metadata` is now a required semantic dependency for `validate_business_rules()`.
|
|
- Date format validation remains, but allowed formats are defined locally in `train_model_params.py`.
|
|
|
|
### 2.4 Input loading pipeline changes in workflow
|
|
|
|
Before:
|
|
- `validate_train_params` directly consumed workflow input.
|
|
|
|
Now:
|
|
1. `load_model_metadata` runs first (fetches model index/schema from plugin store and injects `model_metadata`).
|
|
2. `validate_train_params` runs with enriched payload.
|
|
|
|
This means IO preprocessing now depends on plugin-store metadata resolution before final validation.
|
|
|
|
---
|
|
|
|
## 3) Output contract diff (before vs now)
|
|
|
|
### 3.1 Workflow return (`train_model.run`)
|
|
|
|
Before (`main`):
|
|
- Return type: `None`
|
|
- Workflow side effects were persisted mainly via DB status updates and MLflow artifacts.
|
|
|
|
Now:
|
|
- Return type: `dict[str, Any] | None`
|
|
- Workflow returns the training activity summary when successful.
|
|
|
|
### 3.2 Activity-level training result payload
|
|
|
|
Before (from `Training.train_model` in `main` path):
|
|
- Returned minimal dict:
|
|
- `run_name`
|
|
- `run_dir`
|
|
|
|
Now:
|
|
- Returns extended dict:
|
|
- `run_name`
|
|
- `experiment_name`
|
|
- `run_id`
|
|
- `run_dir`
|
|
|
|
### 3.3 Persistence map by destination (DB, MLflow, MinIO, local filesystem)
|
|
|
|
This section maps where each artifact/metadata goes, in which format, and how that changed from `main`.
|
|
|
|
#### 3.3.1 PostgreSQL (`experiment_run` table)
|
|
|
|
## Before (`main`)
|
|
|
|
- Update path: `update_experiment_run` activity with `UpdateType.MODEL_SAVED`.
|
|
- Persisted on success:
|
|
- `status` transition to `TRAINING_SUCCESS`
|
|
- `run_name` (MLflow run identifier used by current implementation)
|
|
- Persisted on failures:
|
|
- `status` transition to validation/training error statuses
|
|
- `error_message`
|
|
|
|
## Now (current branch)
|
|
|
|
- Same update path and status/error behavior.
|
|
- Even though train activity now returns more metadata (`run_id`, `experiment_name`), current workflow update for `MODEL_SAVED` still forwards mainly `run_name`.
|
|
- Practical effect:
|
|
- DB remains status-centric and run-name-centric
|
|
- richer identifiers exist in workflow return payload, not fully mirrored to DB columns in current flow
|
|
|
|
#### 3.3.2 MLflow (tracking server/artifact store)
|
|
|
|
## Before (`main`)
|
|
|
|
- Persistence orchestration lived in `ModelRepository.save_model()` + `_save_run()`.
|
|
- Typical persisted content:
|
|
- model params (many legacy params such as lags, limits, scaler config, removed intervals)
|
|
- regression metrics (`MSE`, `R2`, `MAE`)
|
|
- model objects:
|
|
- `data_model`
|
|
- `prediction_model`
|
|
- artifacts:
|
|
- `report.html`
|
|
- `train_data.csv`
|
|
- `test_data.csv`
|
|
- optional `model_equation.json`
|
|
- Run naming:
|
|
- computed by querying existing runs and appending sequence (`<experiment>-<n>` style)
|
|
|
|
## Now (current branch)
|
|
|
|
- Persistence orchestrated in `Training._persist_training_artifacts()` and MLflow run context is opened by `SientiaMLflowRepository.start_run(...)`.
|
|
- Persisted content now:
|
|
- model wrapper itself via `wrapper.store_model(name=train_params.model_name)`
|
|
- regression metrics also logged as MLflow params via `mlflow.log_param(...)`:
|
|
- `mse_val`
|
|
- `mae_val`
|
|
- `r2_val`
|
|
- artifacts explicitly logged with `mlflow.log_artifact(...)`:
|
|
- `report.html`
|
|
- `train_data.csv`
|
|
- `test_data.csv`
|
|
- metrics are computed before save (`mse_val`, `mae_val`, `r2_val`) and persisted in the run as params
|
|
- Run identifiers now exposed back to workflow:
|
|
- `experiment_name`
|
|
- `run_name`
|
|
- `run_id`
|
|
- Notable behavioral change:
|
|
- `wrapper._input_example` is cleared (`None`) before storing model.
|
|
|
|
#### 3.3.3 MinIO object storage
|
|
|
|
## Before (`main`)
|
|
|
|
- Read path:
|
|
- single source object downloaded via `StorageRepository.fetch_file(bucket_name, file_name)`
|
|
- Write path:
|
|
- training workflow did not write generated outputs to MinIO in this code path
|
|
- generated artifacts were persisted to MLflow, not uploaded back to MinIO
|
|
- Location:
|
|
- source data in input bucket/key provided by workflow input (`bucket_name` + `file_name`)
|
|
|
|
## Now (current branch)
|
|
|
|
- Read path migrated to `MinioRepository.download_file(...)`.
|
|
- Supports two input objects:
|
|
- mandatory training object: `bucket_name` + `file_name`
|
|
- optional validation object: same `bucket_name` + `val_file_name`
|
|
- Write path:
|
|
- still no artifact upload to MinIO in this workflow path
|
|
- report/CSV outputs continue to flow to MLflow artifacts
|
|
- Location details:
|
|
- bucket resolved from payload (`bucket_name`)
|
|
- object key exactly from payload (`file_name`, optional `val_file_name`)
|
|
- default bucket in env/config is `MINIO_DEFAULT_BUCKET`, but runtime payload can override via `bucket_name`
|
|
|
|
#### 3.3.4 Local filesystem (ephemeral runtime workspace)
|
|
|
|
## Before (`main`)
|
|
|
|
- Temporary run dir created under reports root using run name + timestamp suffix.
|
|
- Artifacts generated locally in that directory:
|
|
- `report.html`
|
|
- `train_data.csv`
|
|
- `test_data.csv`
|
|
- optional `model_equation.json`
|
|
- After MLflow logging, cleanup activity removed temp directory.
|
|
|
|
## Now (current branch)
|
|
|
|
- Temporary run dir managed by `DataManagerRepository` under runtime reports root (`.../reports/temp/<run_name>`).
|
|
- Same artifact family generated locally:
|
|
- `report.html`
|
|
- `train_data.csv`
|
|
- `test_data.csv`
|
|
- optional `model_equation.json` (for `linear_regression`)
|
|
- Cleanup behavior is now tolerant:
|
|
- cleanup runs in guarded `finally`
|
|
- training success is not reverted if cleanup later fails
|
|
|
|
#### 3.3.5 Quick matrix (before vs now)
|
|
|
|
- **Postgres**
|
|
- before: status + run_name + errors
|
|
- now: same persisted shape; workflow return contains extra IDs
|
|
- **MLflow**
|
|
- before: legacy model objects + params/metrics + report/data artifacts
|
|
- now: wrapper-based model persistence + `mse_val`/`mae_val`/`r2_val` as params + report/data artifacts + run_id exposed
|
|
- **MinIO**
|
|
- before: reads 1 CSV input object
|
|
- now: reads 1 or 2 CSV input objects (train + optional validation), still no output upload
|
|
- **Local temp**
|
|
- before: generated artifacts, then cleanup
|
|
- now: generated artifacts, then best-effort cleanup (non-blocking for success result)
|
|
|
|
### 3.4 Cleanup behavior impact on output semantics
|
|
|
|
Before:
|
|
- Cleanup was called directly after training result; failures propagated straightforwardly.
|
|
|
|
Now:
|
|
- Cleanup is in a guarded `finally`.
|
|
- If training succeeded but cleanup fails, workflow warns and does not rollback success semantics.
|
|
- Effective output semantics: successful training result can be returned even if temp cleanup fails.
|
|
|
|
---
|
|
|
|
## 4) Detailed field mapping (old -> new)
|
|
|
|
## Kept (or equivalent role)
|
|
|
|
- `experiment_run_id` -> kept (broader accepted types: int or numeric string)
|
|
- `variable_columns` -> kept
|
|
- `target_variable` -> kept
|
|
- `bucket_name` -> kept
|
|
- `file_name` -> kept
|
|
- `line_separator` -> kept
|
|
- `decimal_separator` -> kept
|
|
- `date_column` -> required (snake_case key; must exist in CSV)
|
|
- `date_format` -> optional in payload; omitted/null/blank resolves to default `yyyy-MM-dd HH:mm:ss`
|
|
- `train_size` -> kept
|
|
- `shuffle` -> kept
|
|
- `model_name` -> kept (now less coupled to legacy model enum)
|
|
|
|
## Added
|
|
|
|
- `model_type` (primary selector for plugin wrapper/index lookup)
|
|
- `data_model_kwargs`
|
|
- `model_kwargs`
|
|
- `opt_params`
|
|
- `val_file_name` (optional second dataset input)
|
|
- `model_id` (optional metadata)
|
|
- `model_metadata` (loaded/required for schema validation)
|
|
- `random_state` (explicit split reproducibility control)
|
|
|
|
## Removed from new required contract
|
|
|
|
- `lag_train`, `lag_val`
|
|
- `rem_static_win`, `static_threshold`
|
|
- `low_lim`, `upp_lim`
|
|
- `window`
|
|
- `use_scaler`, `include_ar`
|
|
- `degree`, `interaction_only`
|
|
- `nan_treatment`
|
|
- `start_date`, `end_date`
|
|
- `scaler_name`
|
|
- `removed_intervals`
|
|
- `support_filters`
|
|
- `experiment_name` (no longer required as top-level client input)
|
|
|
|
---
|
|
|
|
## 5) Internal architecture update notes
|
|
|
|
### 5.1 Repository layer redesign
|
|
|
|
Removed:
|
|
- `model_manager/utils/repository/model_repository.py`
|
|
- `model_manager/utils/repository/training_repository.py`
|
|
- `model_manager/utils/repository/storage_repository.py`
|
|
|
|
Added:
|
|
- `model_manager/utils/repository/data_manager_repository.py`
|
|
|
|
Interpretation:
|
|
- Data preprocessing/report/metrics responsibilities were consolidated into `DataManagerRepository`.
|
|
- Training/model persistence shifted to wrapper + plugin store + MLflow repository integrations.
|
|
|
|
### 5.2 Model engine abstraction migration
|
|
|
|
Before:
|
|
- Strong coupling to local classes in `model_manager.sientia.models` and custom preprocessing/model objects in `TrainModelResult`.
|
|
|
|
Now:
|
|
- Training uses `SientiaModel` wrapper dynamically obtained by `plugin_store.get_model(model_type=...)`.
|
|
- Contract is wrapper-driven (`train`, `transform`, `predict`, `store_model`).
|
|
- The codebase removed `model_manager/sientia/models.py`, `model_serving.py`, and `utils.py`, indicating full migration to externalized model runtime abstraction.
|
|
|
|
### 5.3 Worker/runtime architecture changes
|
|
|
|
- New `prepare_worker.py` centralizes worker setup and autoscaling parameters.
|
|
- Queue names are now runtime-derived:
|
|
- `train_model-<runtime>-queue`
|
|
- `cleanup_files-<runtime>-queue`
|
|
- `worker.py` now installs runtime via plugin store (`plugin_store.install_runtime(runtime_name=...)`) before starting workers.
|
|
- This introduces environment/runtime-aware deployment and model packaging behavior.
|
|
|
|
### 5.4 Synchronous activity and tracking adjustments
|
|
|
|
- `experiment_tracking` migrated from async postgres helper to sync postgres client path (`postgres_sync`).
|
|
- Several activities switched to sync method signatures.
|
|
- Error handling in workflow and DB status update paths is more defensive (secondary failures while persisting error status are logged and do not mask primary failure cause).
|
|
|
|
### 5.5 `TrainModelResult` shape update
|
|
|
|
Before:
|
|
- Stored classic split artifacts (`x_train`, `x_test`, `y_train`, `y_test`) + concrete preprocessing/model objects (`process_data`, `regr`, `scaler_dict`).
|
|
|
|
Now:
|
|
- Stores `train_data`, `val_data` and prediction DataFrames, plus tracking identifiers (`experiment_name`, `run_id`).
|
|
- Result object is less tied to internal estimator classes and more aligned with serializable workflow/model-store integration.
|
|
|
|
---
|
|
|
|
## 6) Net IO compatibility assessment
|
|
|
|
## Input compatibility
|
|
|
|
Not backward compatible with old payloads without adaptation.
|
|
|
|
Key reasons:
|
|
- Legacy required fields removed/ignored by new path.
|
|
- New required fields introduced (`model_type`, `*_kwargs` dicts, runtime metadata flow dependency).
|
|
- Validation pipeline now expects model metadata semantics.
|
|
|
|
## Output compatibility
|
|
|
|
Behavior changed:
|
|
- Workflow now returns a result object (previously `None`).
|
|
- Training summary includes `experiment_name` and `run_id` in addition to `run_name` and `run_dir`.
|
|
- DB update still centered on `run_name`; callers relying only on DB may not see all new output info unless workflow return is consumed.
|
|
|
|
---
|
|
|
|
## 7) Practical migration guidance (client side)
|
|
|
|
To call `train_model` in this branch:
|
|
|
|
1. Send snake_case payload aligned to new `TrainModelParams`.
|
|
2. Always provide:
|
|
- `model_name` slugified model name (ex.: `test_model_name or test-model-name`)
|
|
- `model_type`
|
|
- `data_model_kwargs` (dict)
|
|
- `model_kwargs` (dict)
|
|
- `opt_params` (dict)
|
|
3. Keep `experiment_run_id` numeric (int or numeric string).
|
|
4. Use runtime queue naming consistent with worker runtime:
|
|
- `train_model-<runtime>-queue`
|
|
5. If you need explicit validation split file, send `val_file_name`; otherwise split uses `train_size`/`shuffle`/`random_state`.
|
|
|
|
---
|
|
|
|
## 8) Source references used for this document
|
|
|
|
Primary diffs:
|
|
- `model_manager/workflows/train_model.py`
|
|
- `model_manager/utils/models/train_model_params.py`
|
|
- `model_manager/utils/models/train_model_result.py`
|
|
- `model_manager/activities/training.py`
|
|
- `model_manager/activities/activities.py`
|
|
- `model_manager/activities/experiment_tracking.py`
|
|
- `model_manager/utils/repository/data_manager_repository.py`
|
|
- `model_manager/utils/repository/model_repository.py` (removed)
|
|
- `model_manager/utils/repository/training_repository.py` (removed)
|
|
- `model_manager/utils/repository/storage_repository.py` (removed)
|
|
- `model_manager/worker/worker.py`
|
|
- `model_manager/worker/prepare_worker.py`
|
|
- `README.md`
|
|
- `input-sample.md`
|
|
- `scripts/run_training_test.py`
|
|
|