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:
+        This activity orchestrates the ML training pipeline:
-        1. Validates input parameters
+        1. Validate input parameters.
-        2. Trains the model using TrainingRepository
+        2. Train the model via TrainingRepository.
-        3. Performs post-training calculations
+        3. Perform post-training calculations.
 
         Args:
-            input_data: Configuration for model training operation
+            input_data: Training configuration containing:
-                Required keys:
+                - metadata (dict): Workflow execution metadata.
-                    - metadata (dict): Workflow execution metadata
+                - uploaded_file (BytesIO): Training data already downloaded from MinIO.
-                    - uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
+                - train_params (TrainModelParams | dict): Training parameters.
-                    - train_params (TrainModelParams): Training parameters object
 
         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
+            ValueError: If input validation fails.
-            Exception: If training fails (after sending notification)
+            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(...)
         """
         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:
+        Raises:
-            SientiaMlException if unable to search runs
+            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:
+        Raises:
-            with model_serving.save_experiment(run_name="my_run") as run:
+            Exception: If starting or ending the MLflow run fails.
-                model_serving.log_param("param1", value1)
-                model_serving.log_metric("metric1", value2)
-                # Run is automatically closed here, even if an exception occurs
         """
         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_URL: MLFlow server hostname (default: http://localhost:5080)
-        MLFLOW_PORT: MLFlow server port (default: 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:

todo-list.txt

@@ -1,3 +1,4 @@
 - remover os testes das classes alteradas e refazer de novo depois
-- alterar todos os comentários dos métodos que foram alterados
 - no final alterar o readme
+- verificar erro de logs duplicados
+- Remover do log do metadata o model name e id