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/activities/experiment_tracking.py

@@ -110,6 +110,17 @@ class ExperimentTracking(Postgres):
                 pass
 
     async def _execute_update(self, query: str, params: Mapping[str, Any]) -> dict[str, Any]:
+        """
+        Execute an UPDATE SQL statement asynchronously.
+
+        Args:
+            query: Parameterized SQL string to execute.
+            params: Mapping of parameters for the SQL query.
+
+        Returns:
+            dict: A dictionary containing the affected row count: {'rowcount': int}.
+        """
+
         def _run() -> dict[str, Any]:
             with self.engine.begin() as connection:
                 result = connection.execute(text(query), params)

model_manager/activities/training.py

@@ -113,32 +113,23 @@ class Training(BaseActivity):
         """
         Train a machine learning model.
 
-        This activity orchestrates the complete ML training pipeline:
-        1. Validates input parameters
-        2. Trains the model using TrainingRepository
-        3. Performs post-training calculations
+        This activity orchestrates the ML training pipeline:
+        1. Validate input parameters.
+        2. Train the model via TrainingRepository.
+        3. Perform post-training calculations.
 
         Args:
-            input_data: Configuration for model training operation
-                Required keys:
-                    - metadata (dict): Workflow execution metadata
-                    - uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
-                    - train_params (TrainModelParams): Training parameters object
+            input_data: Training configuration containing:
+                - metadata (dict): Workflow execution metadata.
+                - uploaded_file (BytesIO): Training data already downloaded from MinIO.
+                - train_params (TrainModelParams | dict): Training parameters.
 
         Returns:
-            TrainModelResult: Training result with model, metrics, and data
+            dict: Keys `run_name` and `run_dir` when training and saving succeed.
 
         Raises:
-            ValueError: If input validation fails
-            Exception: If training fails (after sending notification)
-
-        Example:
-            result = await train_model({
-                'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
-                'uploaded_file': BytesIO(csv_data),
-                'train_params': TrainModelParams(...)
-            })
-            # Returns: TrainModelResult(...)
+            ValueError: If input validation fails.
+            Exception: If training fails (after sending notification).
         """
         metadata = input_data.get('metadata', {})
         train_params = input_data['train_params']
@@ -195,6 +186,19 @@ class Training(BaseActivity):
 
     @activity.defn(name='cleanup_resources')
     async def cleanup_resources(self, input_data: dict[str, Any]) -> None:
+        """
+        Cleanup temporary resources created during training.
+
+        Args:
+            input_data: Cleanup configuration containing:
+                - metadata (dict): Workflow execution metadata.
+                - run_dir (str): Temporary directory to remove.
+                - bucket_name (str): MinIO bucket of the uploaded file.
+                - file_name (str): MinIO object key to delete.
+
+        Raises:
+            Exception: If cleanup fails (after sending notification).
+        """
         metadata = input_data.get('metadata', {})
         run_dir = input_data.get('run_dir', '')
         bucket_name = input_data.get('bucket_name', '')

model_manager/sientia/model_serving.py

@@ -66,8 +66,8 @@ class ModelServing:
         Returns:
             pandas.DataFrame: A DataFrame containing run information.
 
-        Raise:
-            SientiaMlException if unable to search runs
+        Raises:
+            SientiaMlException: If unable to search runs.
         """
         try:
             runs = mlflow.search_runs(experiment_names=experiment_names, order_by=order_by)
@@ -82,6 +82,9 @@ class ModelServing:
 
         Args:
             experiment_identifier (str): name or id of the experiment to be setted
+
+        Raises:
+            Exception: If setting the experiment fails.
         """
         mlflow.set_experiment(experiment_identifier)
 
@@ -99,6 +102,9 @@ class ModelServing:
         Security Warning:
             The GitHub token is hardcoded. Consider moving to environment variable
             or using a secure secret management solution (e.g., K8s secrets).
+
+        Raises:
+            Exception: If logging the model fails.
         """
         mlflow.sklearn.log_model(
             sk_model,
@@ -117,6 +123,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the parameter fails.
         """
         mlflow.log_param(key, value)
 
@@ -130,6 +139,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the metric fails.
         """
         mlflow.log_metric(key, value)
 
@@ -146,6 +158,9 @@ class ModelServing:
 
         Returns:
             None
+
+        Raises:
+            Exception: If logging the artifact fails.
         """
         mlflow.log_artifact(local_path=local_path, artifact_path=artifact_path, run_id=run_id)
 
@@ -178,11 +193,8 @@ class ModelServing:
         Yields:
             ActiveRun: object that acts as a context manager wrapping the run's state.
 
-        Example:
-            with model_serving.save_experiment(run_name="my_run") as run:
-                model_serving.log_param("param1", value1)
-                model_serving.log_metric("metric1", value2)
-                # Run is automatically closed here, even if an exception occurs
+        Raises:
+            Exception: If starting or ending the MLflow run fails.
         """
         run = mlflow.start_run(
             run_id=run_id,

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: