From 98c216733b6bf4a5768d17c4803a1259a0071cff Mon Sep 17 00:00:00 2001
From: Bruno Domingues <bruno.domingues@aignosi.com.br>
Date: Wed, 15 Oct 2025 16:19:15 -0300
Subject: [PATCH] SIENTIAPDE-1253: Enhance error logging in train_model
 workflow with rich context for debugging.

---
 model_manager/workflows/train_model.py | 52 ++++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 4 deletions(-)

diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py
index 69aa7c5..29eb632 100644
--- a/model_manager/workflows/train_model.py
+++ b/model_manager/workflows/train_model.py
@@ -238,9 +238,18 @@ class TrainModel:
         except Exception as e:
             # Validation failed: Update status with error
             error_message = str(e)
+            error_type = type(e).__name__
 
+            # Log with rich context for debugging
             workflow.logger.error(
-                f'Validation failed for experiment {experiment_run_id}: {error_message}'
+                f'[VALIDATION_ERROR] Experiment {experiment_run_id} validation failed',
+                extra={
+                    'step': 'validate_training_parameters',
+                    'experiment_run_id': experiment_run_id,
+                    'error_type': error_type,
+                    'error_message': error_message,
+                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
+                },
             )
 
             await self._update_experiment_run(
@@ -322,9 +331,21 @@ class TrainModel:
         except Exception as e:
             # Download or training failed: Update status with error
             error_message = str(e)
+            error_type = type(e).__name__
 
+            # Log with rich context for debugging
             workflow.logger.error(
-                f'Download or training failed for experiment {experiment_run_id}: {error_message}'
+                f'[TRAINING_ERROR] Experiment {experiment_run_id} training failed',
+                extra={
+                    'step': 'download_and_train_model',
+                    'experiment_run_id': experiment_run_id,
+                    'experiment_name': train_params.experiment_name,
+                    'bucket_name': train_params.bucket_name,
+                    'file_name': train_params.file_name,
+                    'error_type': error_type,
+                    'error_message': error_message,
+                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
+                },
             )
 
             await self._update_experiment_run(
@@ -396,9 +417,20 @@ class TrainModel:
         except Exception as e:
             # Model saving failed: Update status with error
             error_message = str(e)
+            error_type = type(e).__name__
 
+            # Log with rich context for debugging
             workflow.logger.error(
-                f'Model saving failed for experiment {experiment_run_id}: {error_message}'
+                f'[MLFLOW_ERROR] Experiment {experiment_run_id} model save failed',
+                extra={
+                    'step': 'save_model_to_mlflow',
+                    'experiment_run_id': experiment_run_id,
+                    'experiment_name': train_result.params.experiment_name,
+                    'run_dir': train_result.run_dir if hasattr(train_result, 'run_dir') else None,
+                    'error_type': error_type,
+                    'error_message': error_message,
+                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
+                },
             )
 
             await self._update_experiment_run(
@@ -477,9 +509,21 @@ class TrainModel:
         except Exception as e:
             # Cleanup failed: Update status with error
             error_message = str(e)
+            error_type = type(e).__name__
 
+            # Log with rich context for debugging
             workflow.logger.error(
-                f'Cleanup failed for experiment {experiment_run_id}: {error_message}'
+                f'[CLEANUP_ERROR] Experiment {experiment_run_id} cleanup failed',
+                extra={
+                    'step': 'cleanup_resources',
+                    'experiment_run_id': experiment_run_id,
+                    'bucket_name': saved_result.params.bucket_name,
+                    'file_name': saved_result.params.file_name,
+                    'run_dir': saved_result.run_dir if hasattr(saved_result, 'run_dir') else None,
+                    'error_type': error_type,
+                    'error_message': error_message,
+                    'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
+                },
             )
 
             await self._update_experiment_run(