plan-mlflowIntegration

Prompt file imported from AnshBajpai05/Hydroscope---Refinery-H2-Optimization (.github/prompts/plan-mlflowIntegration.prompt.md). Copyright stays with the author.

Plan: FCHCU-First MLflow Integration Blueprint

Integrate MLflow in parallel with your existing file-based experiment tracker (no breaking changes), starting in FCHCU and designing reusable utilities so DHDS/DHT and future generation pipelines can adopt the same pattern. The plan logs metrics, params, artifacts, and model versions to MLflow while keeping current registry/state JSON flows intact until deprecation is explicitly approved.

Steps

1. Phase 0: Baseline Mapping and Contract Freeze

2. Confirm and document the current run contract from experiment orchestration in consumption/FCHCU/scripts/experiment_tracking.py (step names, run_id format, summary.json locations, run_state schema, registry update semantics).

3. Define canonical MLflow naming contract: experiment name = project/domain + exp_id, run name = current run_name, tags for pipeline_step, split_method, target_col, project=FCHCU, domain=consumption. This blocks all downstream logging work.

4. Phase 1: Shared Tracking Abstraction (Reusable for consumption + generation)

5. Add a reusable tracking helper module under consumption/FCHCU/scripts (designed to be portable) that wraps MLflow operations and exposes:

6. start_parent_run(context)

7. log_step_metrics(step_name, metrics_dict)

8. log_step_params(step_name, params_dict)

9. log_artifacts(step_name, file_or_dir_paths)

10. register_model_if_eligible(model_uri, model_name, metadata)

11. The helper must support backend-agnostic configuration via environment variables (MLFLOW_TRACKING_URI, MLFLOW_EXPERIMENT_PREFIX, artifact location options), so local file backend and central server are both supported without code changes.

12. Include “safe logging” behavior (skip/mark missing artifacts instead of failing training pipeline).

13. Phase 2: FCHCU Orchestrator Integration (Parent/Child run hierarchy)

14. In consumption/FCHCU/scripts/experiment_tracking.py, initialize one parent MLflow run at pipeline start and keep existing run_state.json/timeline.log updates unchanged.

15. For each pipeline stage (fetching, preparation, preprocess train/test, feature selection, training, deployment, monitoring), create child runs or step-tagged logging blocks and map existing summary extraction into MLflow metrics.

16. Mirror current timeline and status semantics to MLflow tags (status, started_at, completed_at, step failure reason).

17. Keep _finalize_registry behavior; add MLflow final summary tags/metrics there to preserve parity with registry.json.

18. Phase 3: Script-Level Logging Hooks (FCHCU first)

19. In consumption/FCHCU/scripts/model_training.py:

20. Log tuning params (n_trials, cv config), best_params, and best CV/test metrics.

21. Log model artifacts and report outputs (model_artifact.json, summary.json, plots, HTML report, XAI artifacts).

22. Log dataset profile counts (train/test rows, features used, selected feature count).

23. In consumption/FCHCU/scripts/deploy_model.py:

24. Log deploy_artifact.json and production pipeline artifacts.

25. Emit deployment metadata tags (deployment_id, source_experiment, model compatibility version).

26. In consumption/FCHCU/scripts/post_deployment_monitor.py:

27. Log monitoring summary metrics (RMSE/MAE/R2/MAPE/MDA) and drift indicators (PSI/KS aggregates + per-feature artifacts).

28. Log generated monitoring plots and HTML for run-to-run drift visibility in MLflow UI.

29. In consumption/FCHCU/scripts/export_predictions.py:

30. Optionally log exported report metadata (output file path, selected experiment) for traceability.

31. Phase 4: Model Registry Integration (requested for this phase)

32. Define model naming policy for registry (e.g., consumption_fchcu_targetname).

33. Register trained model from FCHCU training/deployment output with version tags that point to exp_id/run_name and key metrics.

34. Define stage transition policy (None/Staging/Production) and minimal promotion gate based on existing test metrics + monitoring thresholds.

35. Preserve current production artifact generation (production_pipeline.dill/.pkl/.json/.onnx) while also creating MLflow-registered versions.

36. Phase 5: Template Rollout to DHDS and DHT (parallelizable after FCHCU validation)

37. Replicate integration pattern to:

38. consumption/DHDS/scripts/experiment_tracking.py

39. consumption/DHDS/scripts/model_training.py

40. consumption/DHDS/scripts/deploy_model.py

41. consumption/DHDS/scripts/post_deployment_monitor.py

42. consumption/DHT/scripts/experiment_tracking.py

43. consumption/DHT/scripts/model_training.py

44. consumption/DHT/scripts/deploy_model.py

45. consumption/DHT/scripts/post_deployment_monitor.py

46. Keep project-specific experiment prefixes/tags so cross-project comparison works in one MLflow instance.

47. Phase 6: Future-Ready Generation Scaffolding (no generation code changes yet)

48. Publish a folder-agnostic integration guide (required tags, run hierarchy, artifact conventions, helper usage) so new generation models can plug in immediately.

49. Define a minimal onboarding checklist for any new project folder: configure tracking URI, register experiment prefix, map summary.json to metric schema.

50. Phase 7: Verification and Acceptance

51. Execute one full FCHCU run and verify:

52. Parent + step-level logs present in MLflow.

53. Metrics parity between summary.json and MLflow values.

54. Artifacts visible/downloadable in MLflow.

55. Model version created in MLflow Model Registry with correct tags.

56. Execute one monitoring run and verify drift metrics/plots are logged.

57. Regression-check existing local outputs are still produced (registry.json, run_state.json, timeline.log, report files).

58. Validate backend portability by switching MLFLOW_TRACKING_URI between local and server-compatible value.

59. Phase 8: Retraining Actuation, Trigger Policy, and Visualization

60. Add a retraining decision engine in consumption/FCHCU/scripts/post_deployment_monitor.py that evaluates trigger classes and emits a machine-readable decision artifact.

61. Define trigger classes with project-configurable thresholds:

62. Data drift trigger (for example PSI > 0.2 on key features for N consecutive windows).

63. Concept drift trigger (for example rolling RMSE degradation > 15 percent versus production baseline for N windows).

64. Absolute performance trigger (for example R2 below threshold or MAPE above threshold).

65. Freshness trigger (maximum model age or maximum days since last successful retrain).

66. Business/manual trigger (operator or business event override).

67. Implement debounce and cooldown controls to avoid retraining loops:

68. Consecutive-window confirmation before actuation.

69. Cooldown period after successful retrain unless severe trigger is active.

70. Support governance modes:

71. Auto mode: trigger-approved retraining launches automatically.

72. Human-in-loop mode: candidate retrain is logged and awaits explicit approval.

73. Emit decision states and transition timestamps:

74. no_action, candidate, approved, executed, promoted, rejected, failed.

75. Add an orchestrated retrain launch path in consumption/FCHCU/scripts/experiment_tracking.py that links monitoring run lineage to retraining run lineage.

76. Persist and log decision evidence as artifacts:

77. retrain_decision.json

78. trigger_evidence_summary.json

79. gate_report.json

80. Log explicit MLflow tags and metrics for visibility:

81. Tags: trigger_reason, decision_state, source_model_name, source_model_version, retrain_policy_version, governance_mode.

82. Metrics: trigger_data_drift, trigger_concept_drift, trigger_perf_drop, trigger_staleness, severity_score, actuation_latency_minutes, promotion_gate_pass.

83. Define visualization requirements in MLflow UI (or companion dashboard backed by MLflow runs):

84. Retraining Decision Board with one row per monitoring cycle.

85. Trigger Timeline with markers for candidate, approved, executed, promoted/rejected.

86. Trigger Attribution chart by reason code and top drift-contributing features.

87. Lineage panel linking monitor run to retrain run to model registry version.

88. SLA panel: mean time to detect, mean time to retrain, mean time to promote, retrain success rate.

Relevant files

• <REPO_ROOT>/consumption/FCHCU/scripts/experiment_tracking.py — parent orchestration, run lifecycle, registry finalization, step summary ingestion.
• <REPO_ROOT>/consumption/FCHCU/scripts/model_training.py — tuning/training metrics and artifact emission.
• <REPO_ROOT>/consumption/FCHCU/scripts/deploy_model.py — deployment artifact generation and production bundle metadata.
• <REPO_ROOT>/consumption/FCHCU/scripts/post_deployment_monitor.py — post-deployment performance and drift outputs.
• <REPO_ROOT>/consumption/FCHCU/scripts/export_predictions.py — optional trace logging for exported outputs.
• <REPO_ROOT>/consumption/FCHCU/data/experiment/registry.json — current registry retained in parallel mode.
• <REPO_ROOT>/consumption/DHDS/scripts/experiment_tracking.py — rollout target after FCHCU validation.
• <REPO_ROOT>/consumption/DHT/scripts/experiment_tracking.py — rollout target after FCHCU validation.

Verification

1. Run FCHCU orchestrator end-to-end and confirm MLflow run tree, metrics, artifacts, and tags.
2. Compare key metrics from MLflow vs data/experiment/<exp_id>/gold/reports/*/summary.json for exact parity on selected metrics (RMSE, MAE, R2, MAPE).
3. Confirm Model Registry version creation and metadata tagging from training/deploy outputs.
4. Run post_deployment_monitor.py and verify drift metrics plus visual artifacts are attached to the corresponding MLflow run.
5. Repeat one smoke run with alternate tracking URI configuration to verify backend-agnostic behavior.
6. Force each trigger class in a controlled test and verify decision_state transitions, cooldown behavior, and governance routing (auto vs human-in-loop).
7. Validate retrain lineage links from monitoring run to retraining run to registered model version.
8. Validate actuation latency and decision metrics are queryable and chartable in MLflow.

Decisions

• In scope: FCHCU-first implementation plan, reusable design for future folders, parallel mode (keep current tracking), include MLflow Model Registry.
• In scope: explicit retraining trigger policy, automated/manual actuation modes, and retraining visualization via MLflow-compatible run metadata.
• Out of scope for initial implementation: deleting registry.json/run_state.json, redesigning ProductionPipeline serialization format, generation-folder code edits (folder is currently empty).
• Assumption: Existing summary.json outputs remain the source of truth for compatibility during migration; MLflow mirrors them.

Further Considerations

1. Registry gate policy recommendation: promote to Staging when test RMSE improves over current Production by configurable threshold and monitoring drift is below threshold in the latest run.
2. Schema standardization recommendation: define a fixed metric/tag key set now to avoid dashboard fragmentation across DHDS/DHT/FCHCU/generation.