feat: enhance training and experiment tracking functionality

- Updated `Activities` class to improve garbage collection handling.
- Enhanced error messaging in `ExperimentTracking` for better clarity on update failures.
- Refactored `Training` class to streamline exception handling and improve type hints.
- Introduced new methods in `TrainModelParams` for better handling of experiment run IDs and model metadata.
- Added functionality to extract model equations in `DataManagerRepository` for linear regression models.

model_manager/workflows/cleanup_files.py

@@ -13,7 +13,7 @@ with workflow.unsafe.imports_passed_through():
     from typing import Any
 
     from model_manager.activities.activities import Activities
-    from model_manager.workflows.train_model import POD_ID, no_retry_policy
+    from model_manager.workflows.train_model import no_retry_policy
 
     TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
 
@@ -45,7 +45,7 @@ class CleanupFiles:
         # Metadata for tracking
         metadata = {
             'metadata': {
-                'pod_id': POD_ID,
+                'pod_id': os.getenv('POD_ID'),
                 'workflow_name': 'cleanup_files',
             }
         }

model_manager/workflows/train_model.py

@@ -23,8 +23,9 @@ with workflow.unsafe.imports_passed_through():
     from model_manager.utils.models.experiment_status import ExperimentStatus
     from model_manager.utils.models.train_model_params import TrainModelParams
 
-    # Activity Timeouts (in seconds) - Configurable via environment variables
-    # Defaults are designed to handle large files (up to 200MB)
+    # Activity timeouts (seconds). Tune per environment (large uploads, long training).
+    # Training uses no_retry_policy: extend TIMEOUT_TRAIN_MODEL instead of adding retries
+    # to avoid duplicate MLflow side effects. Cleanup/delete uses network_retry_policy.
     TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
     TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '2700'))
     TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '120'))
@@ -71,7 +72,7 @@ class TrainModel:
     """
 
     @workflow.run
-    async def run(self, input_data: dict[str, Any]) -> dict[str, str | None] | None:
+    async def run(self, input_data: dict[str, Any]) -> dict[str, Any] | None:
         """
         Execute the complete model training workflow.
 
@@ -95,9 +96,11 @@ class TrainModel:
             ValueError: If experiment_run_id is missing or invalid
         """
         experiment_run_id = self._validate_experiment_run_id(input_data)
+        input_data = {**input_data, 'experiment_run_id': experiment_run_id}
+
         model_name = input_data.get('model_name')
         model_id = input_data.get('model_id')
-        
+
         metadata = {
             'metadata': {
                 'experiment_run_id': experiment_run_id,
@@ -112,7 +115,7 @@ class TrainModel:
         )
 
         training_succeeded = False
-        train_result: dict[str, str | None] | None = None
+        train_result: dict[str, Any] | None = None
 
         try:
             train_result = await self._train_model(
@@ -128,9 +131,15 @@ class TrainModel:
                         run_dir=train_result.get('run_dir'),
                         metadata=metadata,
                     )
+                else:
+                    pass
             except Exception:
-                if training_succeeded:
-                    raise
+                # If cleanup fails after training failed, there is nothing extra to log (DB not committed).
+                if training_succeeded:  # pragma: no branch
+                    workflow.logger.warning(
+                        'cleanup_resources failed after successful training; model and DB status '
+                        'are already committed. Temp files may remain until scheduled cleanup.',
+                    )
         return train_result
 
     def _validate_experiment_run_id(self, input_data: dict[str, Any]) -> int:
@@ -227,8 +236,11 @@ class TrainModel:
                     status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
                     error_message=self._extract_error_message(e),
                 )
-            except Exception:
-                pass
+            except Exception as secondary:
+                workflow.logger.warning(
+                    'Failed to persist ORCHESTRATOR_VALIDATION_ERROR to experiment_run: %s',
+                    secondary,
+                )
             raise
 
     async def _train_model(
@@ -236,7 +248,7 @@ class TrainModel:
         train_params: TrainModelParams,
         experiment_run_id: int,
         metadata: dict[str, Any],
-    ) -> dict[str, str | None]:
+    ) -> dict[str, Any]:
         """
         Download file from MinIO and train model.
 
@@ -250,7 +262,7 @@ class TrainModel:
             metadata: Workflow execution metadata
 
         Returns:
-            TrainModelResult: Training result from train_model activity
+            dict[str, Any]: Serializable training summary from the train_model activity
 
         Raises:
             Exception: If download or training fails (after updating DB status)
@@ -285,8 +297,11 @@ class TrainModel:
                     status=ExperimentStatus.TRAINING_ERROR,
                     error_message=self._extract_error_message(e),
                 )
-            except Exception:
-                pass
+            except Exception as secondary:
+                workflow.logger.warning(
+                    'Failed to persist TRAINING_ERROR status to experiment_run: %s',
+                    secondary,
+                )
             raise
 
     async def _cleanup_resources(