sientia-dataops-model-manager/docs/train-model-workflow-io-diff-main-vs-current-branch.md

# 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` -> kept optional
- `date_format` -> kept optional
- `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`