SIENTIAPDE-1241: Refactor: Improve documentation, exception handling, and configuration in model manager. This commit enhances clarity and robustness by adding detailed docstrings to methods, standardizing exception handling with custom types, and simplifying MLflow configuration.

model_manager/utils/connectors_config.py

@@ -42,8 +42,7 @@ def build_mlflow_config() -> dict[str, Any]:
     It handles server connection and authentication parameters.
 
     Environment Variables:
-        MLFLOW_HOST: MLFlow server hostname (default: http://localhost)
-        MLFLOW_PORT: MLFlow server port (default: 5080)
+        MLFLOW_URL: MLFlow server hostname (default: http://localhost:5080)
         MLFLOW_USERNAME: MLFlow username (default: aignosi)
         MLFLOW_PASSWORD: MLFlow password (default: aignosi)
 

model_manager/utils/exceptions.py

@@ -1,5 +1,31 @@
+"""
+Custom exception types for the Model Manager.
+
+This module defines domain-specific exceptions used across the training
+workflow to convey additional context (e.g., flags indicating which steps
+completed successfully) without altering control flow semantics.
+"""
+
+
 class ModelTrainingError(Exception):
+    """
+    Exception raised when the model training workflow fails.
+
+    This exception carries flags indicating whether the model was trained
+    and/or saved successfully, enabling the workflow to map errors to
+    appropriate experiment statuses.
+    """
+
     def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None):
+        """
+        Initialize ModelTrainingError with training state flags.
+
+        Args:
+            model_trained: True if the training step completed successfully.
+            model_saved: True if the model saving step completed successfully.
+            message: Optional custom error message. If None, a default message
+                including the state flags is generated.
+        """
         self.model_trained = model_trained
         self.model_saved = model_saved
 

model_manager/utils/models/experiment_status.py

@@ -14,6 +14,7 @@ class ExperimentStatus(str, Enum):
     to maintain compatibility with existing database records and monitoring systems.
 
     Attributes:
+        ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation.
         MAGE_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
         TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
         TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.

model_manager/utils/models/train_model_params.py

@@ -78,14 +78,6 @@ class TrainModelParams:
             ValueError: If any required field is missing or None
             TypeError: If any field has an incorrect type
             KeyError: If any required key is missing from the dictionary
-
-        Example:
-            >>> input_data = {
-            ...     'variable_columns': ['var1', 'var2'],
-            ...     'lag_train': 5,
-            ...     # ... other fields
-            ... }
-            >>> params = TrainModelParams.from_dict(input_data)
         """
         return cls(
             variable_columns=cls._check_none(
@@ -179,10 +171,6 @@ class TrainModelParams:
 
         Raises:
             ValueError: If any business rule is violated
-
-        Example:
-            >>> params = TrainModelParams.from_dict(data)
-            >>> params.validate_business_rules()  # Raises ValueError if invalid
         """
         # Validate train_size range (10-100%)
         if not 10 <= self.train_size <= 100: