SIENTIAPDE-1231

Enhance README and repository utilities for clarity and functionality

- Updated README.md to improve descriptions and structure, adding detailed sections for features, workflows, and architecture.
- Enhanced MinioRepository with comprehensive docstrings for methods and class attributes, improving usability and documentation.
- Refined MLFlowRepository with clearer method descriptions and improved logging for better observability and maintainability.

README.md

@@ -1,42 +1,108 @@
 # Sientia DataOps Laborious
 
-A high-performance, scalable machine learning prediction system built on Temporal.io for industrial data processing and ML model inference. The Laborious system provides enterprise-grade ML model management, batch prediction processing, and real-time data export capabilities with comprehensive data quality validation and monitoring.
+A comprehensive, Temporal-based ML orchestration system for industrial data processing and model inference. Laborious delivers enterprise-grade batch prediction, model management, optional real-time export (OPC), and automated retraining with strong data quality validation and observability.
+
+## 📑 Table of Contents
+
+- [Features](#features)
+  - [Core Functionality](#core-functionality)
+  - [Advanced Capabilities](#advanced-capabilities)
+  - [Development & Quality Assurance](#development--quality-assurance)
+- [Architecture](#architecture)
+  - [Architecture Principles](#architecture-principles)
+  - [Key Components](#key-components)
+  - [Data Flow Architecture](#data-flow-architecture)
+  - [Security Architecture](#security-architecture)
+- [Workflows](#workflows)
+  - [Predictions Batch Workflow](#1-predictions-batch-workflow-predictions_batchpy)
+  - [Prediction Process Workflow](#2-prediction-process-workflow-prediction_processpy)
+  - [Format and Export Prediction Workflow](#3-format-and-export-prediction-workflow-format_and_export_predictionpy)
+  - [Minimal Retrain Workflow](#4-minimal-retrain-workflow-minimal_retrainpy)
+- [Installation & Setup](#installation--setup)
+  - [Prerequisites](#prerequisites)
+  - [Environment Setup](#environment-setup)
+  - [Temporal Namespace Setup](#temporal-namespace-setup)
+  - [Local Development Setup](#local-development-setup)
+- [How to Run](#how-to-run)
+  - [Running the Laborious Application](#running-the-laborious-application)
+  - [Running Tests and Coverage](#running-tests-and-coverage)
+  - [Manual Test Execution](#manual-test-execution)
+  - [Manual Application Execution](#manual-application-execution)
+- [Code Quality & Validation](#code-quality--validation)
+  - [Overview](#overview)
+  - [Validation Tools](#validation-tools)
+  - [Tools Installation](#tools-installation)
+  - [Complete Validation](#complete-validation)
+  - [Automatic Fixes](#automatic-fixes)
+  - [Configuration](#configuration)
+  - [CI/CD Integration](#cicd-integration)
+  - [Best Practices](#best-practices)
+- [Testing](#testing)
+  - [Test Structure](#test-structure)
+  - [Test Execution](#test-execution)
+- [Monitoring and Metrics](#monitoring-and-metrics)
+  - [Application Health Metrics](#application-health-metrics)
+  - [Prediction Operation Metrics](#prediction-operation-metrics)
+  - [OPC Export Metrics](#opc-export-metrics)
+  - [Data Quality Metrics](#data-quality-metrics)
+- [Configuration](#configuration-1)
+  - [Environment Variables](#environment-variables)
+  - [OPC Configuration](#opc-configuration)
+  - [Workflow Configuration](#workflow-configuration)
+- [Development](#development)
+  - [Project Structure](#project-structure)
+  - [Adding New Features](#adding-new-features)
+- [Troubleshooting](#troubleshooting)
+  - [Common Issues](#common-issues)
+  - [Debug Mode](#debug-mode)
+- [Performance Tuning](#performance-tuning)
+  - [Key Parameters](#key-parameters)
+  - [Scaling Considerations](#scaling-considerations)
+- [Contributing](#contributing)
+  - [Code Quality Standards](#code-quality-standards)
+- [License](#license)
+- [Support](#support)
 
 ## Features
 
 ### Core Functionality
-- **Batch Prediction Processing**: High-throughput ML model inference using MLFlow models
+- **Batch Prediction Processing**: High-throughput ML inference using MLFlow models
-- **Temporal Workflow Orchestration**: Robust workflow management with automatic retry policies and fault tolerance
+- **Temporal Workflow Orchestration**: Robust workflow management with retries and fault tolerance
-- **Data Quality Gates**: Configurable filtering for data validation, MLFlow API responses, and custom validation rules
+- **Data Quality Gates**: Configurable filtering for input data and MLFlow API responses
-- **Multi-Model Support**: Flexible ML model management with retention policies and versioning
+- **Multi-Model Support**: Flexible model management with retention and versioning
-- **Real-time Data Export**: PostgreSQL persistence and OPC server integration for industrial systems
+- **Optional Real-time Export**: PostgreSQL persistence and OPC server integration for industrial systems
-- **Comprehensive Monitoring**: Prometheus metrics and detailed logging for operational visibility
+- **Comprehensive Monitoring**: Prometheus metrics and structured logging for observability
 
 ### Advanced Capabilities
-- **Incremental Data Processing**: Timestamp-based data loading to avoid reprocessing
+- **Incremental Data Processing**: Timestamp-based loading to avoid reprocessing
 - **Configurable Data Retention**: Model retention policies with automatic cleanup
-- **Notification System**: Integrated alerting and notification management via MongoDB
+- **Notification System**: Integrated alerting via MongoDB
-- **Scalable Architecture**: Kubernetes-ready deployment with horizontal scaling support
+- **Scalable Architecture**: Kubernetes-ready with horizontal scaling
-- **Model Retraining**: Automated model retraining workflows with production model updates
+- **Model Retraining**: Automated retraining workflows with production model updates
+
+### Development & Quality Assurance
+- **Code Quality Tools**: Ruff (lint/format), mypy (types), Bandit (security)
+- **Automated Validation**: `validate.sh` and CI quality gates
+- **Comprehensive Testing**: pytest with async support and high coverage
+- **Type Safety**: Static type checking with mypy
+- **Coverage Visualization**: Coverage Gutters integration
 
 ## Architecture
 
-The Laborious system uses a Temporal-based workflow architecture with clear separation of concerns and robust error handling. The architecture is designed for high availability, scalability, and operational excellence in production ML environments.
+Laborious uses a Temporal-based architecture with strong separation of concerns and defensive error handling for production ML.
-
 
 ### Architecture Principles
 
 #### 1. **Separation of Concerns**
-- **Worker Layer**: Manages Temporal workers, task queues, and application lifecycle
+- **Worker Layer**: Temporal workers, task queues, lifecycle
-- **Workflow Layer**: Orchestrates business logic and process coordination
+- **Workflow Layer**: Business orchestration and coordination
-- **Activity Layer**: Implements specific operations and external system interactions
+- **Activity Layer**: External system interactions and isolated operations
-- **Data Layer**: Handles data persistence, caching, and external service connections
+- **Data Layer**: Persistence, caching, connectors
 
 #### 2. **Fault Tolerance & Resilience**
-- **Automatic Retry Policies**: Configurable retry strategies for transient failures
+- **Automatic Retry Policies** for transient failures
-- **Circuit Breaker Pattern**: Prevents cascading failures in external service calls
+- **Graceful Degradation** and circuit breaking for dependencies
-- **Graceful Degradation**: System continues operating with reduced functionality
+- **Detailed Error Handling** with notifications
-- **Comprehensive Error Handling**: Detailed error reporting and notification integration
 
 #### 3. **Scalability & Performance**
 - **Horizontal Scaling**: Multiple worker instances for load distribution
@@ -53,101 +119,57 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa
 ### Key Components
 
 #### **Worker (`laborious/worker/worker.py`)**
-- **Purpose**: Main application orchestrator managing Temporal workers and task queues
+- Temporal client setup, worker lifecycle, task queues
-- **Responsibilities**:
+- Metrics server initialization, notification handler setup
-  - Temporal client initialization and connection management
+- Graceful shutdown and autoscaling-friendly behavior
-  - Worker lifecycle management and graceful shutdown
-  - Task queue configuration and load balancing
-  - Prometheus metrics server initialization
-  - Notification handler setup and configuration
-  - OPC server connection management
-- **Key Features**:
-  - Automatic scaling with `PollerBehaviorAutoscaling`
-  - Health check endpoints for Kubernetes liveness/readiness probes
-  - Graceful shutdown with cleanup procedures
-  - Multi-instance deployment support
-  - Two dedicated task queues: `predictions_batch-queue` and `minimal_retrain-queue`
 
 #### **Workflows (`laborious/workflows/`)**
-- **PredictionsBatch**: Main entry point for batch prediction pipelines
+- `predictions_batch.py`: Batch prediction entry point
-- **PredictionProcess**: Core prediction pipeline with MLFlow integration
+- `sub_workflows/prediction_process.py`: Core prediction pipeline
-- **FormatAndExportPrediction**: Data formatting and export operations
+- `sub_workflows/format_and_export_prediction.py`: Formatting and export
-- **MinimalRetrain**: Automated model retraining and deployment
+- `minimal_retrain.py`: Automated model retraining and production update
-- **Key Features**:
-  - Temporal workflow definitions with retry policies
-  - Child workflow orchestration and delegation
-  - Comprehensive error handling and recovery
-  - Configurable timeout and retry strategies
 
 #### **Activities (`laborious/activities/`)**
-- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance
+- `gates.py`: Data quality validation and filtering
-- **Gates**: Data quality validation and filtering mechanisms
+- `mlflow.py`: Transform and predict operations
-- **MLFlow**: Model transformation and prediction operations
+- `opc.py`: OPC UA export to industrial systems (optional)
-- **OPC**: Real-time data export to industrial OPC servers
+- `activities.py`: Aggregates activity interfaces
-- **Key Features**:
-  - Multiple inheritance pattern for unified activity interface
-  - Configurable filter policies and validation rules
-  - MLFlow model serving integration with configurable flavors
-  - OPC UA client with certificate-based authentication
-  - Comprehensive error handling and notification integration
-  - Support for multiple OPC servers with independent configurations
 
 #### **Data Services (`laborious/utils/`)**
-- **Connectors Config**: Environment variable-based configuration management
+- `connectors_config.py`: Env-driven configuration builders
-- **Repository**: Data access layer for MLFlow and OPC operations
+- `repository/model_repository.py`: MLFlow operations and retraining
-  - `model_repository.py`: MLFlow model operations and retraining
+- `repository/opc_repository.py`: OPC communication and writes
-  - `opc_repository.py`: OPC server communication and data writing
+- `filters/conditional_filters.py` and `filters/mlflow_filters.py`
-- **Filters**: Data quality validation and MLFlow response filtering
-  - `conditional_filters.py`: Input data validation filters
-  - `mlflow_filters.py`: MLFlow API response validation filters
-- **Key Features**:
-  - Environment variable-based configuration with sensible defaults
-  - Connection pool management and optimization
-  - Security credential management
-  - Configuration validation and error handling
-  - Support for multiple OPC servers and MLFlow model flavors
 
 ### Data Flow Architecture
 
 #### **1. Batch Prediction Pipeline**
 ```
-Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform → 
+Input Data (PostgreSQL) → Data Quality Gates → MLFlow Transform →
-MLFlow Prediction → Response Validation → Export (PostgreSQL + OPC)
+MLFlow Prediction → Response Validation → Export (PostgreSQL [+ OPC])
 ```
 
 #### **2. Model Retraining Pipeline**
 ```
-Training Data → Model Retraining → Quality Validation → 
+Training Data → Model Retraining → Quality Validation →
 Production Update → Notification & Monitoring
 ```
 
-#### **3. Real-time Export Pipeline**
-```
-Prediction Results → Data Formatting → OPC Server Write → 
-Success/Failure Metrics → Notification System
-```
-
 ### Security Architecture
 
 #### **Authentication & Authorization**
-- **Certificate-based OPC Authentication**: Secure industrial communication
+- **MLFlow API Authentication**: Username/password
-- **MLFlow API Authentication**: Username/password with secure transmission
+- **Database Security**: Encrypted connections and credential management
-- **Database Connection Security**: Encrypted connections with credential management
+- **OPC Certificates** (if enabled): Client/server certs
-- **Kubernetes Secrets Integration**: Secure credential storage and access
+- **Kubernetes Secrets**: Secure secret storage
 
 #### **Network Security**
-- **TLS/SSL Encryption**: Secure communication channels
+- TLS/SSL, network policies, service mesh, firewalls, VPN
-- **Network Isolation**: Kubernetes network policies and service mesh
-- **Firewall Rules**: Controlled access to external services
-- **VPN Integration**: Secure remote access and management
 
 #### **Data Security**
-- **Data Encryption**: At-rest and in-transit encryption
+- At-rest/in-transit encryption, RBAC, audit logging, lifecycle management
-- **Access Control**: Role-based access control (RBAC)
-- **Audit Logging**: Comprehensive access and operation logging
-- **Data Retention**: Configurable data lifecycle management
 
-## 🔄 Workflows
+## Workflows
 
 ### 1. Predictions Batch Workflow (`predictions_batch.py`)
 
@@ -351,8 +373,9 @@ flowchart LR
 - Temporal server/cluster
 - PostgreSQL database
 - MLFlow server
-- OPC server(s)
+- MinIO object storage (for MLFlow artifacts)
 - MongoDB server (for notifications)
+- OPC server(s) if using OPC export
 
 **Note**: External dependencies must be available either through:
 - Kubernetes cluster deployment

laborious/utils/repository/minio_repository.py

@@ -1,3 +1,11 @@
+"""
+MinIO repository utilities.
+
+This module provides a lightweight repository around a MinIO/S3-compatible
+object storage using boto3. It supports creating buckets on demand and
+storing/loading pandas DataFrames in Parquet format.
+"""
+
 from io import BytesIO
 from typing import Any
 
@@ -10,6 +18,21 @@ from sientia_do.observability.logger import Logger
 
 
 class MinioRepository:
+    """
+    Repository for interacting with a MinIO (S3-compatible) object storage.
+
+    This class encapsulates a reusable `boto3` S3 client and convenience
+    helpers to persist and retrieve pandas DataFrames as Parquet files.
+
+    Attributes:
+        storage_options (dict): Options compatible with pandas s3fs usage.
+        minio_bucket (str): Default bucket name used for operations.
+        minio_endpoint_url (str): MinIO endpoint URL.
+        minio_region_name (str): MinIO region name.
+        s3_client (Any): Reusable S3 client from `boto3`.
+        logger (Logger): Observability logger.
+        notification_handler (NotificationHandler): Notifications handler.
+    """
     def __init__(
         self,
         minio_endpoint_url: str,
@@ -20,6 +43,17 @@ class MinioRepository:
         logger: Logger,
         notification_handler: NotificationHandler,
     ):
+        """Initialize the repository and S3 client.
+
+        Args:
+            minio_endpoint_url (str): MinIO endpoint URL.
+            minio_access_key (str): Access key (AK).
+            minio_secret_key (str): Secret key (SK).
+            minio_region_name (str): Region name for the client.
+            minio_default_bucket (str): Default bucket name to operate on.
+            logger (Logger): Logger instance for structured logs.
+            notification_handler (NotificationHandler): Notification handler.
+        """
         # MinIO settings shared with pandas s3fs
         self.storage_options = {
             'key': minio_access_key,
@@ -54,11 +88,14 @@ class MinioRepository:
         self.notification_handler = notification_handler
 
     def close(self):
+        """Close the underlying S3 client."""
         self.s3_client.close()
 
     def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None:
-        """
+        """Ensure the default bucket exists; create it if missing.
-        Ensure the MinIO bucket exists; create it if necessary.
+
+        Args:
+            metadata (dict[str, Any]): Metadata used for structured logging.
         """
 
         try:
@@ -71,6 +108,14 @@ class MinioRepository:
     def store_dataframe_as_parquet(
         self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any]
     ):
+        """Persist a DataFrame as a Parquet object in the default bucket.
+
+        Args:
+            dataframe (DataFrame): DataFrame to persist.
+            uri (str): Human-friendly URI used for logging context.
+            object_name (str): Object key (path/key within the bucket).
+            metadata (dict[str, Any]): Metadata used for structured logging.
+        """
         self.ensure_bucket_exists(metadata)
 
         self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata)
@@ -83,6 +128,15 @@ class MinioRepository:
         self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata)
 
     def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame:
+        """Load a Parquet object from the default bucket into a DataFrame.
+
+        Args:
+            object_key (str): Object key to retrieve from the bucket.
+            metadata (dict[str, Any]): Metadata used for structured logging.
+
+        Returns:
+            DataFrame: Loaded DataFrame.
+        """
         self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata)
 
         response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)

laborious/utils/repository/model_repository.py

@@ -1,13 +1,13 @@
 """
-MLFlow Repository
+MLflow repository utilities
 
-This module contains the MLFlowRepository class,
+This module provides the `MLFlowRepository` class and helpers to interact with
-which is responsible for handling the communication with MLFlow tracking server.
+an MLflow tracking server and model registry. It covers model discovery,
+downloading/loading with multiple flavors, cached operations with retention
+policies, transformation/prediction interfaces, retraining workflows, and
+production model promotion.
 
-It includes methods for model management, caching, retraining, and serving operations
+Capabilities:
-using MLFlow's tracking and model registry capabilities.
-
-The repository provides comprehensive functionality for:
 - Model loading and caching with retention policies
 - Data transformation and prediction operations
 - Model retraining workflows
@@ -37,6 +37,15 @@ INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'
 
 
 def force_memory_release(logger: Logger):
+    """Attempt to release memory from the Python process.
+
+    Executes a garbage collection cycle and calls `malloc_trim(0)` on glibc
+    where available to return free memory to the OS. This may be a no-op on
+    non-glibc systems.
+
+    Args:
+        logger (Logger): Logger for observability.
+    """
     gc.collect()
 
     try:
@@ -48,6 +57,14 @@ def force_memory_release(logger: Logger):
 
 class MLFlowRepository:
     def __init__(self, host: str, username: str, password: str, logger: Logger):
+        """Initialize MLflow client and base state.
+
+        Args:
+            host (str): MLflow tracking URI.
+            username (str): MLflow username.
+            password (str): MLflow password.
+            logger (Logger): Logger instance.
+        """
         # set tracking uri
         mlflow.set_tracking_uri(host)
 
@@ -64,15 +81,15 @@ class MLFlowRepository:
     """
 
     def get_model_uri(self, run_id: str, prediction: bool = True):
-        """
+        """Build the artifact URI for a run's model.
-        Get the model URI based on the run_id.
 
         Args:
-            run_id (str): The run_id of the model.
+            run_id (str): MLflow run identifier.
-            prediction (bool): Whether to get prediction model URI (default: True)
+            prediction (bool): If True, return `prediction_model` URI,
+                otherwise return `data_model` URI.
 
         Returns:
-            str: The model URI.
+            str: Artifact URI to the selected model within the run.
         """
         run_info = mlflow.get_run(run_id)
         if prediction:
@@ -82,15 +99,14 @@ class MLFlowRepository:
         return model_uri
 
     def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str:
-        """
+        """Resolve the run_id for a registered model at a given stage.
-        Get the run_id of a model based on its name and stage.
 
         Args:
-            model_name (str): The name of the model.
+            model_name (str): Registered model name.
-            stage (str): The stage of the model.
+            stage (str): Desired stage (e.g., 'Production').
 
         Returns:
-            str: The run_id of the model.
+            str: Run ID for the latest version at the given stage.
         """
         # Use search_registered_models instead of deprecated get_latest_versions
         registered_models = self.client.search_registered_models(
@@ -120,14 +136,14 @@ class MLFlowRepository:
 
     def get_next_run_name(self, model_name: str) -> str:
         """
-        Generate the next run name for a specific MLFlow model.
+        Generate the next run name for a specific MLflow model.
 
         This method calculates the next sequential run number for a model
         by searching existing runs and incrementing the count. It ensures
         unique run names for model training and retraining operations.
 
         Args:
-            model_name (str): The name of the MLFlow model
+            model_name (str): The name of the MLflow model
 
         Returns:
             str: The next run name in format 'model_name-run_number'
@@ -140,20 +156,20 @@ class MLFlowRepository:
         self, experiment_name: str, create_if_not_exists: bool = False
     ) -> Experiment:
         """
-        Retrieve MLFlow experiment ID by experiment name.
+        Retrieve MLflow experiment by name, optionally creating it.
 
         This method searches for an MLFlow experiment by name and
         returns its unique identifier. It provides error handling
         for non-existent experiments.
 
         Args:
-            experiment_name (str): Name of the MLFlow experiment
+            experiment_name (str): Name of the MLflow experiment
 
         Returns:
-            int: MLFlow experiment ID
+            Experiment: MLflow experiment object
 
         Raises:
-            ValueError: If the experiment name is not found
+            ValueError: If the experiment name is not found and creation is disabled
         """
         experiment = mlflow.get_experiment_by_name(experiment_name)
 
@@ -166,7 +182,14 @@ class MLFlowRepository:
         return experiment
 
     def get_model_params(self, run_id: str):
-        """Obtém os parâmetros de uma run"""
+        """Fetch parameters associated with a given MLflow run.
+
+        Args:
+            run_id (str): Run identifier to inspect.
+
+        Returns:
+            dict: Mapping of parameter names to values.
+        """
         run_info = mlflow.get_run(run_id)
         return run_info.data.params
 
@@ -176,14 +199,14 @@ class MLFlowRepository:
 
     def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str:
         """
-        Downloads artifacts from a specific MLFlow run.
+        Download artifacts from the latest production run of a model.
 
         Args:
-            model_name (str): Name of the model
+            model_name (str): Registered model name.
-            artifact_path (str): Path to the artifact within the run
+            artifact_path (str): Relative path to artifacts within the run.
 
         Returns:
-            str: Path to the downloaded artifacts
+            str: Local filesystem path where artifacts are saved.
         """
         run_id = self.get_model_run_id(model_name=model_name, stage='Production')
         output_dir = f'{ARTIFACTS_PATH}/{model_name}'
@@ -201,7 +224,7 @@ class MLFlowRepository:
 
     def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any:
         """
-        Downloads a predictive model from the MLflow Model Registry.
+        Load a predictive model from the MLflow Model Registry.
 
         Args:
             model_name (str): The name of the model to download from the registry.
@@ -212,7 +235,7 @@ class MLFlowRepository:
             mlflow.pyfunc.PyFuncModel: The loaded predictive model.
 
         Notes:
-            - The model is fetched from the "production" stage of the MLflow Model Registry.
+            - The model is fetched from the "Production" stage of the MLflow Model Registry.
             - Warnings during the model loading process are suppressed.
         """
         model_uri = f'models:/{model_name}/production'
@@ -230,7 +253,7 @@ class MLFlowRepository:
 
     def load_transform_model(self, model_name: str, flavor: str) -> Any:
         """
-        Downloads the latest production version of a specified transformation model.
+        Load the latest Production version of a transformation model.
 
         This method retrieves the latest production model run ID for the given
         model name, constructs the model URI, and loads the model using MLflow.
@@ -241,11 +264,11 @@ class MLFlowRepository:
             artifact_path (str | None): Path to compressed artifacts if model is compressed
 
         Returns:
-            Any: The loaded model object, as returned by `mlflow.sklearn.load_model`.
+            Any: The loaded model object, depending on the flavor used.
 
         Raises:
             Exception: If the model run ID or URI cannot be retrieved, or if the
-                       model cannot be loaded.
+                model cannot be loaded.
         """
 
         latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
@@ -266,7 +289,7 @@ class MLFlowRepository:
         self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False
     ) -> tuple[Any, str | None]:
         """
-        Download model based on type (predict or transform).
+        Download model based on type ("predict" or "transform").
 
         Args:
             model_name (str): Name of the model to download
@@ -275,7 +298,7 @@ class MLFlowRepository:
             load_wrapper (bool): Whether to load wrapper
 
         Returns:
-            tuple[Any, str]: Model object and artifact path if model is compressed
+            tuple[Any, str | None]: Model object and optional artifact path.
         """
 
         self.logger.info(
@@ -317,17 +340,19 @@ class MLFlowRepository:
 
     def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
         """
-        Detect and parse datetime index from data. index must be a timestamp like column.
+        Normalize DataFrame index to the expected timestamp string format.
-        This function must detect the timestamp type (pandas Timestamp or datetime) and convert it to DATETIME_FORMAT_WITH_TZ.
+
-        If the index is a string, must be in format DATETIME_FORMAT_WITH_TZ.
+        The index must be timestamp-like. If the index is:
-        If another type or format, must raise an error.
+        - string: it must match `DATETIME_FORMAT_WITH_TZ`
+        - datetime or pandas Timestamp: it will be converted to that format
+        Any other type raises a ValueError.
 
         Args:
-            data (pd.DataFrame): DataFrame with timestamp index
+            data (pd.DataFrame): DataFrame with timestamp index.
-            metadata (dict): Metadata for logging
+            metadata (dict): Metadata for structured logging.
 
         Returns:
-            pd.DataFrame: DataFrame with converted datetime index
+            pd.DataFrame: DataFrame with converted datetime index.
         """
         index = data.index
 
@@ -368,14 +393,14 @@ class MLFlowRepository:
 
     def check_cache_retention(self, cache: dict, retention: int) -> bool:
         """
-        Check if cache is still valid based on retention time.
+        Check whether cached model data is still valid.
 
         Args:
-            cache (dict): Cached model data
+            cache (dict): Cached model data with a 'timestamp' key.
-            retention (int): Retention time in minutes
+            retention (int): Retention time in minutes.
 
         Returns:
-            bool: True if cache is still valid, False if expired
+            bool: True if cache is still valid, False if expired.
         """
         current_time = datetime.now()
         cache_time = cache['timestamp']
@@ -385,17 +410,14 @@ class MLFlowRepository:
 
     def handle_valid_model(self, model_name: str, cache: dict) -> dict:
         """
-        Handle valid cached model by returning appropriate model configuration.
+        Return the cached model configuration when retention is valid.
 
         Args:
-            model_name (str): Name of the model
+            model_name (str): Name of the model (for logging/consistency).
-            model_type (str): Type of model ('predict' or 'transform')
+            cache (dict): Cached model data structure.
-            compressed (bool): Whether model is compressed
-            retention_target (str): Retention target ('model' or 'artifact')
-            cache (dict): Cached model data
 
         Returns:
-            dict: Model configuration with model and artifact path
+            dict: Model configuration.
         """
         self.logger.debug(f'Model {model_name} is still valid, using cached version')
 
@@ -406,8 +428,8 @@ class MLFlowRepository:
         Clean up outdated cached model and its artifacts.
 
         Args:
-            model_name (str): Name of the model
+            model_name (str): Name of the model.
-            model_key (str): Cache key for the model
+            model_key (str): Cache key for the model.
 
         Returns:
             None
@@ -419,16 +441,16 @@ class MLFlowRepository:
 
     def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any:
         """
-        Get model with caching support based on retention policy.
+        Retrieve a model with caching support based on retention policy.
 
         Args:
             model_name (str): Name of the model to retrieve
-            retention (int): Cache retention time in minutes (0 = no cache)
+            retention (int): Cache retention time in minutes (0 = no cache).
             model_type (str): Type of model ('predict' or 'transform')
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
 
         Returns:
-            Any: Model object
+            Any: Model object.
         """
         # Retention is 0, download a new model
         if retention <= 0:
@@ -488,16 +510,16 @@ class MLFlowRepository:
         self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str
     ) -> pd.DataFrame | ndarray:
         """
-        Get transformed data using cached transform model.
+        Execute a cached operation using the requested model.
 
         Args:
-            model_name (str): Name of the transform model
+            model_name (str): Registered model name.
-            data (pd.DataFrame): Data to transform
+            data (pd.DataFrame): Input data.
-            retention (int): Cache retention time in minutes
+            retention (int): Cache retention in minutes.
-            flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
+            flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
 
         Returns:
-            pd.DataFrame: Transformed data
+            pd.DataFrame | ndarray: Operation result.
         """
         if operation not in ['transform', 'predict']:
             raise ValueError("Invalid operation. Use 'transform' or 'predict'.")
@@ -531,7 +553,7 @@ class MLFlowRepository:
         target_name: str | None = None,
     ) -> dict[str, dict[str, Any]]:
         """
-        Create a new MLFlow experiment for model retraining.
+        Prepare models and data for a retraining run.
 
         This method sets up the complete environment for model retraining by:
         1. Loading the current production prediction model
@@ -541,18 +563,16 @@ class MLFlowRepository:
         5. Setting up the MLFlow experiment context
 
         Args:
-            model_name (str): Name of the MLFlow model to retrain
+            model_name (str): Name of the MLflow model to retrain.
-            data (pd.DataFrame): Training data for model retraining
+            data (pd.DataFrame): Training data for model retraining.
-            transform_flavor (str): Flavor for transformation model
+            transform_flavor (str): Flavor for transformation model.
-            predict_flavor (str): Flavor for prediction model
+            predict_flavor (str): Flavor for prediction model.
-            target_name (str): Target name
+            target_name (str | None): Optional target column; if None, use model target.
-            metadata (dict): Metadata for logging
+            metadata (dict): Metadata for logging.
 
         Returns:
-            tuple: (prediction_model, data_model, experiment)
+            dict[str, dict[str, Any]]: Mapping with prepared `prediction_model` and
-                - prediction_model: Loaded prediction model for retraining
+                `data_model`, including optional artifact paths.
-                - data_model: Fitted transformation model
-                - experiment: MLFlow experiment name
         """
 
         self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata)
@@ -659,6 +679,14 @@ class MLFlowRepository:
         return retrain_data
 
     def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict):
+        """Log a model into the active MLflow run.
+
+        Args:
+            model_data (dict): Model holder with keys 'model' and optional 'artifact_path'.
+            flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
+            model_type (str): Artifact name, e.g., 'prediction_model' or 'data_model'.
+            metadata (dict): Metadata for structured logging.
+        """
         model = model_data['model']
 
         self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata)
@@ -688,7 +716,7 @@ class MLFlowRepository:
         predict_flavor: str = 'sklearn',
     ) -> dict:
         """
-        Execute the complete model retraining process in MLFlow.
+        Execute the complete model retraining process in MLflow.
 
         This method performs the actual model retraining by:
         1. Starting a new MLFlow run with descriptive metadata
@@ -700,17 +728,15 @@ class MLFlowRepository:
         Args:
             prediction_model: MLFlow prediction model to retrain
             data_model: MLFlow transformation model to retrain
-            experiment (str): MLFlow experiment name for the retraining
+            experiment (str): MLflow experiment name for the retraining.
-            model_name (str): Name of the model being retrained
+            model_name (str): Name of the model being retrained.
-            data (pd.DataFrame): Training data used for retraining
+            data (pd.DataFrame): Training data used for retraining.
-            transform_flavor (str): Flavor for transformation model
+            transform_flavor (str): Flavor for transformation model.
-            predict_flavor (str): Flavor for prediction model
+            predict_flavor (str): Flavor for prediction model.
-            metadata (dict): Metadata for logging
+            metadata (dict): Metadata for logging.
 
         Returns:
-            tuple: (status_message, experiment_name)
+            dict: Metadata about the created run and experiment.
-                - status_message (str): Success confirmation message
-                - experiment_name (str): Name of the experiment
         """
 
         prediction_model = retrain_data['prediction_model']
@@ -795,16 +821,16 @@ class MLFlowRepository:
         self, run_id: str, model_name: str, metadata: dict
     ) -> dict:
         """
-        Update production model with a specific MLFlow run.
+        Promote a specific run's model to Production.
 
         This method promotes a model from a specific MLFlow run to
         production stage. It handles model registration, versioning,
         and stage transitions with proper error handling.
 
         Args:
-            run_id (str): MLFlow run ID containing the model to promote
+            run_id (str): MLflow run ID containing the model to promote.
-            model_name (str): Name of the MLFlow model
+            model_name (str): Name of the MLflow model.
-            metadata (dict): Metadata for logging
+            metadata (dict): Metadata for logging.
 
         Returns:
             dict: Model update metadata containing:
@@ -863,7 +889,7 @@ class MLFlowRepository:
         5. Manages model lifecycle based on retention policy (cleanup artifacts if needed)
 
         Parameters:
-            model_name (str): The name of the MLFlow model to use for transformation.
+            model_name (str): The name of the MLflow model to use for transformation.
             data (pd.DataFrame): The input data to be transformed by the model.
             model_config (dict): Model configuration parameters
             metadata (dict): Metadata for logging
@@ -934,7 +960,7 @@ class MLFlowRepository:
         9. Handles any exceptions and returns structured error information
 
         Parameters:
-            model_name (str): The name of the MLFlow model to use for prediction.
+            model_name (str): The name of the MLflow model to use for prediction.
             data (pd.DataFrame): The input data to make predictions on.
             model_retention (int): Cache retention time in minutes (0 = no caching).
             model_config (dict): Model configuration parameters
@@ -1031,15 +1057,13 @@ class MLFlowRepository:
             data (pd.DataFrame): Training data for model retraining. Must contain
                                 all features required by both transformation and
                                 prediction models, including target variable.
-            model_name (str): Name of the MLFlow model to retrain. Must exist
+            model_name (str): Name of the MLflow model to retrain. Must exist
-                            in the MLFlow Model Registry in Production stage.
+                in the MLflow Model Registry in Production stage.
             model_config (dict): Model configuration parameters
             metadata (dict): Metadata for logging
 
         Returns:
-            tuple: Retraining operation results containing:
+            dict: Retraining operation results and experiment details.
-                - status_message (str): Success confirmation message or error details
-                - experiment_name (str): MLFlow experiment identifier for tracking
 
         Raises:
             mlflow.exceptions.MlflowException: If model not found in registry
@@ -1129,7 +1153,7 @@ class MLFlowRepository:
         3. Returns comprehensive update metadata
 
         Args:
-            experiment (str): MLFlow experiment name containing the retraining runs.
+            experiment (str): MLflow experiment name containing the retraining runs.
                            Must be a valid experiment that exists in MLFlow.
             model_name (str): Name of the MLFlow model to update. Must exist
                             in the MLFlow Model Registry.