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