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', '')