From 26fca36d825d0bd8dcab68da2ff31774f711e4f7 Mon Sep 17 00:00:00 2001 From: Kou-Kinoshita Date: Wed, 29 Oct 2025 16:45:30 -0300 Subject: [PATCH] SIENTIAPDE-1241: Refactored method into smaller ones --- PR_DESCRIPTION.md | 305 ++++++++++++++++++ .../activities/experiment_tracking.py | 158 +++++---- 2 files changed, 401 insertions(+), 62 deletions(-) create mode 100644 PR_DESCRIPTION.md diff --git a/PR_DESCRIPTION.md b/PR_DESCRIPTION.md new file mode 100644 index 0000000..54dd3a0 --- /dev/null +++ b/PR_DESCRIPTION.md @@ -0,0 +1,305 @@ +# Model Equation Artifact - Implementation + +## ๐Ÿ“‹ Summary + +This PR implements automatic extraction and storage of model equations as JSON artifacts in MLflow for Linear Regression models. The equation is extracted after training and logged alongside other artifacts, providing transparency and reproducibility. + +**Issue:** SIENTIAPDE-1321 +**Type:** Feature Enhancement +**Component:** Model Training & MLflow Integration + +--- + +## ๐ŸŽฏ What Changed + +### New Feature: Model Equation Extraction & Storage + +After training a Linear Regression model, the system now: +1. โœ… Extracts coefficients and intercept from the trained model +2. โœ… Generates equation in multiple formats (dict, string, LaTeX) +3. โœ… Saves as `model_equation.json` artifact in MLflow +4. โœ… Logs artifact automatically with other training outputs + +### Example Output +```json +{ + "target_variable": "03CV020/CORRENTE_N_M1_PV(Value)", + "coefficients": { + "303-WIT-200(Value)": 1.234567, + "304-TIT-100(Value)": -0.567890 + }, + "intercept": 10.123456, + "equation_string": "y = 10.123456 + 1.234567 * x1 + -0.567890 * x2", + "latex_equation": "y = 10.123456 + 1.234567 \\cdot x1 + -0.567890 \\cdot x2", + "model_type": "Linear Regression" +} +``` + +--- + +## ๐Ÿ“ Changes by File + +### Core Implementation + +#### `model_manager/utils/models/train_model_result.py` +**Added 2 new fields to TrainModelResult:** +```python +equation: dict | None = None # Model equation metadata +equation_path: str | None = None # Path to equation JSON file +``` + +#### `model_manager/utils/repository/training_repository.py` +**Added equation extraction logic:** +- **New method:** `_extract_model_equation()` - Extracts coefficients, intercept, and generates equation strings +- **Modified:** `after_train_calculation()` - Integrated equation extraction into post-training pipeline + +**Key implementation:** +```python +def _extract_model_equation(self, regr: LinearRegressionModel, params: TrainModelParams) -> dict: + """Extract linear regression equation and generate multiple formats.""" + coefficients = regr.regr.coef_ + intercept = regr.regr.intercept_ + + # Map variables to coefficients + coefficients_dict = { + var: float(coef) + for var, coef in zip(params.variable_columns, coefficients) + } + + # Generate equation strings (plain and LaTeX) + # Return structured dict with all formats +``` + +#### `model_manager/utils/repository/model_repository.py` +**Added artifact generation and logging:** +- **Modified:** `_generate_report()` - Creates `model_equation.json` file when equation data exists +- **Modified:** `_save_run()` - Logs equation artifact to MLflow +- **Added:** `import json` for JSON serialization + +**Key changes:** +```python +# In _generate_report() +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) + +# In _save_run() +if data.equation_path and path.exists(data.equation_path): + self.model_serving.log_artifact(data.equation_path) +``` + +--- + +## ๐Ÿงช Testing + +### Test Coverage Improvements +- **Project Coverage:** 66.27% โ†’ **84.99%** (+18.72%) โœจ +- **Total Tests:** 220 โ†’ **271** (+51 new unit tests) + +### New Test Files + +#### `tests/utils/repository/test_training_repository.py` (NEW) +**23 unit tests covering:** +- โœ… `_extract_model_equation()` - All equation extraction scenarios +- โœ… `_init_data_preprocessor()` - Preprocessor configuration +- โœ… `_init_scaler_dict()` - Scaler initialization (MinMaxScaler, Z_Scaler) +- โœ… `after_train_calculation()` - Post-training pipeline with equation + +**Coverage:** `training_repository.py` went from 17.65% โ†’ **84.31%** (+66.66%) + +#### `tests/utils/repository/test_model_repository.py` (ENHANCED) +**Added 18 new tests:** +- โœ… Equation artifact generation with/without equation data +- โœ… Equation artifact logging to MLflow +- โœ… Error handling for all edge cases +- โœ… Validation of data integrity + +**Coverage:** `model_repository.py` went from 21.43% โ†’ **98.74%** (+77.31%) + +#### `tests/utils/models/test_train_model_result.py` (UPDATED) +**Updated field count validation:** +- Changed expected field count from 17 โ†’ 19 (added `equation` and `equation_path`) + +--- + +## โœ… Quality Validation + +All quality checks passing: + +| Check | Tool | Status | +|-------|------|--------| +| โœ… Code Formatting | `ruff format` | Pass | +| โœ… Linting | `ruff check` | Pass | +| โœ… Type Checking | `mypy` | Pass | +| โœ… Security Analysis | `bandit` | Pass | +| โœ… Unit Tests | `pytest` | 271/271 Pass | +| โœ… Test Coverage | `pytest-cov` | 84.99% | + +--- + +## ๐Ÿš€ Benefits + +### For Data Scientists +- **Transparency:** Easy access to model coefficients without loading the model +- **Documentation:** LaTeX format ready for reports and papers +- **Debugging:** Quick inspection of model structure + +### For MLOps +- **Auditability:** Model equations tracked in MLflow alongside metrics +- **Reproducibility:** Complete mathematical representation stored +- **Integration:** JSON format easily consumable by other systems + +### For Business +- **Explainability:** Clear mathematical relationship between inputs and outputs +- **Compliance:** Model logic explicitly documented +- **Trust:** Transparent model behavior + +--- + +## ๐Ÿ” How to Review + +### 1. Review Data Model Changes +**File:** `model_manager/utils/models/train_model_result.py` +- Check that new fields are optional (`None` defaults) +- Verify field types are correct + +### 2. Review Extraction Logic +**File:** `model_manager/utils/repository/training_repository.py` +- Review `_extract_model_equation()` implementation +- Check coefficient mapping logic +- Verify equation string generation + +### 3. Review Artifact Handling +**File:** `model_manager/utils/repository/model_repository.py` +- Check JSON file creation in `_generate_report()` +- Verify MLflow logging in `_save_run()` +- Confirm error handling + +### 4. Review Test Coverage +**Files:** `tests/utils/repository/test_*.py` +- Verify test scenarios cover edge cases +- Check mock usage is appropriate +- Confirm assertions are meaningful + +### 5. Run Validation +```bash +# Linux/Mac +./validate.sh + +# Windows +ruff format --check . +ruff check . +mypy model_manager/ +pytest tests/ +``` + +--- + +## โš ๏ธ Breaking Changes + +**None** - This is a backward-compatible feature addition. + +- Existing workflows continue to work unchanged +- Equation extraction only activates for Linear Regression models +- All new fields are optional with `None` defaults + +--- + +## ๐Ÿ“Š Test Results + +### End-to-End Testing +โœ… **Validated on Linux (Ubuntu)** - Full training workflow executed successfully +- Training completed +- Equation artifact generated +- JSON logged to MLflow +- Artifact structure validated + +โŒ **Windows testing** - Encountered infrastructure issues (HTTP/2 port-forwarding via Lens) +- Issue is environment-specific, not code-related +- Unit tests pass on Windows + +### Unit Test Summary +``` +====================== 271 passed, 10 warnings in 14.23s ====================== + +Coverage: +- model_manager/utils/repository/training_repository.py 84.31% (+66.66%) +- model_manager/utils/repository/model_repository.py 98.74% (+77.31%) +- TOTAL PROJECT COVERAGE 84.99% (+18.72%) +``` + +--- + +## ๐Ÿ“š Additional Context + +### Architecture Decisions +- **Separation of Concerns:** Equation extraction in `TrainingRepository` (business logic), artifact creation in `ModelRepository` (infrastructure) +- **Optional Implementation:** Won't break if equation extraction fails +- **Multiple Formats:** Supports different use cases (API, documentation, UI) + +### Design Choices +- **Optional fields:** Backward compatible, won't break existing code +- **UTF-8 encoding:** Supports international variable names +- **6 decimal precision:** Balances readability and accuracy +- **JSON format:** Easy integration with other systems + +### Testing Strategy +- **Unit tests:** Validate individual components +- **Integration tests:** Validate complete workflow (tested on Linux) +- **Error handling:** All edge cases covered + +--- + +## ๐Ÿ”ฎ Future Enhancements + +This implementation provides a foundation for: +- Support for other model types (polynomial, tree-based, neural networks) +- Equation visualization in MLflow UI +- Coefficient drift tracking across runs +- Export to PMML/ONNX formats +- Interactive equation explorer + +--- + +## ๐Ÿ“ฆ Files Changed + +### Core Implementation (3 files) +- `model_manager/utils/models/train_model_result.py` +- `model_manager/utils/repository/training_repository.py` +- `model_manager/utils/repository/model_repository.py` + +### Tests (3 files) +- `tests/utils/repository/test_training_repository.py` (NEW - 469 lines) +- `tests/utils/repository/test_model_repository.py` (ENHANCED - 805 lines) +- `tests/utils/models/test_train_model_result.py` (UPDATED) + +### Tooling (1 file) +- `validate.sh` (Fixed pytest argument duplication) + +### Scripts (1 file) +- `scripts/run_training_test.py` (Added noqa comments for security warnings) + +**Total: 8 files modified, 1 new file** + +--- + +## โœ… Checklist + +- [x] Code follows project style guidelines +- [x] All tests passing (271/271) +- [x] Test coverage increased significantly (+18.72%) +- [x] No breaking changes +- [x] Documentation updated +- [x] Type hints added +- [x] Error handling implemented +- [x] Backward compatible +- [x] Security checks passing +- [x] End-to-end tested (Linux) + +--- + +## ๐Ÿ™‹ Questions? + +Feel free to ask questions or request changes. This is a foundational feature that improves model transparency and can be extended based on feedback. + diff --git a/model_manager/activities/experiment_tracking.py b/model_manager/activities/experiment_tracking.py index 911de68..386cad5 100644 --- a/model_manager/activities/experiment_tracking.py +++ b/model_manager/activities/experiment_tracking.py @@ -125,6 +125,99 @@ class ExperimentTracking(Postgres): return await asyncio.to_thread(_run) + def _build_status_update_query( + self, status: str | None, experiment_run_id: int + ) -> tuple[str, dict[str, Any]]: + """Build SQL query for simple status update.""" + if not isinstance(status, str) or not status: + raise ValueError('status is required for STATUS update type') + + sql_query = """ + UPDATE experiment_run + SET status = :status, updated_at = :updated_at + WHERE id = :experiment_run_id + """ + + query_params = { + 'status': status, + 'updated_at': datetime.now(UTC), + 'experiment_run_id': experiment_run_id, + } + + return sql_query, query_params + + def _build_status_with_error_query( + self, status: str | None, error_message: str | None, experiment_run_id: int + ) -> tuple[str, dict[str, Any]]: + """Build SQL query for status update with error message.""" + if not isinstance(status, str) or not status: + raise ValueError('status is required for STATUS_WITH_ERROR update type') + + if not isinstance(error_message, str) or not error_message: + raise ValueError('error_message is required for STATUS_WITH_ERROR update type') + + # Truncate error message if too long + truncated_error = error_message[:1024] if len(error_message) > 1024 else error_message + + sql_query = """ + UPDATE experiment_run + SET status = :status, error_message = :error_message, updated_at = :updated_at + WHERE id = :experiment_run_id + """ + + query_params = { + 'status': status, + 'error_message': truncated_error, + 'updated_at': datetime.now(UTC), + 'experiment_run_id': experiment_run_id, + } + + return sql_query, query_params + + def _build_model_saved_query( + self, run_name: str | None, status: str | None, experiment_run_id: int + ) -> tuple[str, dict[str, Any]]: + """Build SQL query for model saved update.""" + if not isinstance(run_name, str) or not run_name: + raise ValueError('run_name is required for MODEL_SAVED update type') + + if not isinstance(status, str) or not status: + raise ValueError('status is required for MODEL_SAVED update type') + + sql_query = """ + UPDATE experiment_run + SET run_name = :run_name, status = :status, updated_at = :updated_at + WHERE id = :experiment_run_id + """ + + query_params = { + 'run_name': run_name, + 'status': status, + 'updated_at': datetime.now(UTC), + 'experiment_run_id': experiment_run_id, + } + + return sql_query, query_params + + def _get_update_query_and_params( + self, update_type: str, experiment_run_id: int, input_data: dict[str, Any] + ) -> tuple[str, dict[str, Any]]: + """Get SQL query and parameters based on update type.""" + status = input_data.get('status') + error_message = input_data.get('error_message') + run_name = input_data.get('run_name') + + if update_type == UpdateType.STATUS: + return self._build_status_update_query(status, experiment_run_id) + + if update_type == UpdateType.STATUS_WITH_ERROR: + return self._build_status_with_error_query(status, error_message, experiment_run_id) + + if update_type == UpdateType.MODEL_SAVED: + return self._build_model_saved_query(run_name, status, experiment_run_id) + + raise ValueError(f'Invalid update_type: {update_type}') + @activity.defn(name='update_experiment_run') async def update_experiment_run(self, input_data: dict[str, Any]) -> None: """ @@ -153,70 +246,11 @@ class ExperimentTracking(Postgres): experiment_run_id = input_data['experiment_run_id'] update_type = input_data['update_type'] status = input_data.get('status') - error_message = input_data.get('error_message') - run_name = input_data.get('run_name') try: - query_params: dict[str, Any] - - if update_type == UpdateType.STATUS: - if not isinstance(status, str) or not status: - raise ValueError('status is required for STATUS update type') - - sql_query = """ - UPDATE experiment_run - SET status = :status, updated_at = :updated_at - WHERE id = :experiment_run_id - """ - - query_params = { - 'status': status, - 'updated_at': datetime.now(UTC), - 'experiment_run_id': experiment_run_id, - } - elif update_type == UpdateType.STATUS_WITH_ERROR: - if not isinstance(status, str) or not status: - raise ValueError('status is required for STATUS_WITH_ERROR update type') - - if not isinstance(error_message, str) or not error_message: - raise ValueError('error_message is required for STATUS_WITH_ERROR update type') - - if len(error_message) > 1024: - error_message = error_message[:1024] - - sql_query = """ - UPDATE experiment_run - SET status = :status, error_message = :error_message, updated_at = :updated_at - WHERE id = :experiment_run_id - """ - - query_params = { - 'status': status, - 'error_message': error_message, - 'updated_at': datetime.now(UTC), - 'experiment_run_id': experiment_run_id, - } - elif update_type == UpdateType.MODEL_SAVED: - if not isinstance(run_name, str) or not run_name: - raise ValueError('run_name is required for MODEL_SAVED update type') - - if not isinstance(status, str) or not status: - raise ValueError('status is required for MODEL_SAVED update type') - - sql_query = """ - UPDATE experiment_run - SET run_name = :run_name, status = :status, updated_at = :updated_at - WHERE id = :experiment_run_id - """ - - query_params = { - 'run_name': run_name, - 'status': status, - 'updated_at': datetime.now(UTC), - 'experiment_run_id': experiment_run_id, - } - else: - raise ValueError(f'Invalid update_type: {update_type}') + sql_query, query_params = self._get_update_query_and_params( + update_type, experiment_run_id, input_data + ) result = await self._execute_update(sql_query, query_params)