From 644a43093ac6107eaf7f5b2b743e5112e1885f5c Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 16 Oct 2025 11:01:36 -0300 Subject: [PATCH] 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 | 205 ++++++++-------- .../utils/repository/minio_repository.py | 58 ++++- .../utils/repository/model_repository.py | 218 ++++++++++-------- 3 files changed, 291 insertions(+), 190 deletions(-) diff --git a/README.md b/README.md index 500153c..9c1e85f 100644 --- a/README.md +++ b/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 -- **Temporal Workflow Orchestration**: Robust workflow management with automatic retry policies and fault tolerance -- **Data Quality Gates**: Configurable filtering for data validation, MLFlow API responses, and custom validation rules -- **Multi-Model Support**: Flexible ML model management with retention policies and versioning -- **Real-time Data Export**: PostgreSQL persistence and OPC server integration for industrial systems -- **Comprehensive Monitoring**: Prometheus metrics and detailed logging for operational visibility +- **Batch Prediction Processing**: High-throughput ML inference using MLFlow models +- **Temporal Workflow Orchestration**: Robust workflow management with retries and fault tolerance +- **Data Quality Gates**: Configurable filtering for input data and MLFlow API responses +- **Multi-Model Support**: Flexible model management with retention and versioning +- **Optional Real-time Export**: PostgreSQL persistence and OPC server integration for industrial systems +- **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 -- **Scalable Architecture**: Kubernetes-ready deployment with horizontal scaling support -- **Model Retraining**: Automated model retraining workflows with production model updates +- **Notification System**: Integrated alerting via MongoDB +- **Scalable Architecture**: Kubernetes-ready with horizontal scaling +- **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 -- **Workflow Layer**: Orchestrates business logic and process coordination -- **Activity Layer**: Implements specific operations and external system interactions -- **Data Layer**: Handles data persistence, caching, and external service connections +- **Worker Layer**: Temporal workers, task queues, lifecycle +- **Workflow Layer**: Business orchestration and coordination +- **Activity Layer**: External system interactions and isolated operations +- **Data Layer**: Persistence, caching, connectors #### 2. **Fault Tolerance & Resilience** -- **Automatic Retry Policies**: Configurable retry strategies for transient failures -- **Circuit Breaker Pattern**: Prevents cascading failures in external service calls -- **Graceful Degradation**: System continues operating with reduced functionality -- **Comprehensive Error Handling**: Detailed error reporting and notification integration +- **Automatic Retry Policies** for transient failures +- **Graceful Degradation** and circuit breaking for dependencies +- **Detailed Error Handling** with notifications #### 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 -- **Responsibilities**: - - Temporal client initialization and connection management - - 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` +- Temporal client setup, worker lifecycle, task queues +- Metrics server initialization, notification handler setup +- Graceful shutdown and autoscaling-friendly behavior #### **Workflows (`laborious/workflows/`)** -- **PredictionsBatch**: Main entry point for batch prediction pipelines -- **PredictionProcess**: Core prediction pipeline with MLFlow integration -- **FormatAndExportPrediction**: Data formatting and export operations -- **MinimalRetrain**: Automated model retraining and deployment -- **Key Features**: - - Temporal workflow definitions with retry policies - - Child workflow orchestration and delegation - - Comprehensive error handling and recovery - - Configurable timeout and retry strategies +- `predictions_batch.py`: Batch prediction entry point +- `sub_workflows/prediction_process.py`: Core prediction pipeline +- `sub_workflows/format_and_export_prediction.py`: Formatting and export +- `minimal_retrain.py`: Automated model retraining and production update #### **Activities (`laborious/activities/`)** -- **Activities**: Main activity orchestrator combining all functionality through multiple inheritance -- **Gates**: Data quality validation and filtering mechanisms -- **MLFlow**: Model transformation and prediction operations -- **OPC**: Real-time data export to industrial OPC servers -- **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 +- `gates.py`: Data quality validation and filtering +- `mlflow.py`: Transform and predict operations +- `opc.py`: OPC UA export to industrial systems (optional) +- `activities.py`: Aggregates activity interfaces #### **Data Services (`laborious/utils/`)** -- **Connectors Config**: Environment variable-based configuration management -- **Repository**: Data access layer for MLFlow and OPC operations - - `model_repository.py`: MLFlow model operations and retraining - - `opc_repository.py`: OPC server communication and data writing -- **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 +- `connectors_config.py`: Env-driven configuration builders +- `repository/model_repository.py`: MLFlow operations and retraining +- `repository/opc_repository.py`: OPC communication and writes +- `filters/conditional_filters.py` and `filters/mlflow_filters.py` ### Data Flow Architecture #### **1. Batch Prediction Pipeline** ``` -Input Data (PostgreSQL) β†’ Data Quality Gates β†’ MLFlow Transform β†’ -MLFlow Prediction β†’ Response Validation β†’ Export (PostgreSQL + OPC) +Input Data (PostgreSQL) β†’ Data Quality Gates β†’ MLFlow Transform β†’ +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 with secure transmission -- **Database Connection Security**: Encrypted connections with credential management -- **Kubernetes Secrets Integration**: Secure credential storage and access +- **MLFlow API Authentication**: Username/password +- **Database Security**: Encrypted connections and credential management +- **OPC Certificates** (if enabled): Client/server certs +- **Kubernetes Secrets**: Secure secret storage #### **Network Security** -- **TLS/SSL Encryption**: Secure communication channels -- **Network Isolation**: Kubernetes network policies and service mesh -- **Firewall Rules**: Controlled access to external services -- **VPN Integration**: Secure remote access and management +- TLS/SSL, network policies, service mesh, firewalls, VPN #### **Data Security** -- **Data Encryption**: At-rest and in-transit encryption -- **Access Control**: Role-based access control (RBAC) -- **Audit Logging**: Comprehensive access and operation logging -- **Data Retention**: Configurable data lifecycle management +- At-rest/in-transit encryption, RBAC, audit logging, 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 diff --git a/laborious/utils/repository/minio_repository.py b/laborious/utils/repository/minio_repository.py index b38af69..03b0f15 100644 --- a/laborious/utils/repository/minio_repository.py +++ b/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 MinIO bucket exists; create it if necessary. + """Ensure the default bucket exists; create it if missing. + + 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) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 72707c1..3361d9a 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1,13 +1,13 @@ """ -MLFlow Repository +MLflow repository utilities -This module contains the MLFlowRepository class, -which is responsible for handling the communication with MLFlow tracking server. +This module provides the `MLFlowRepository` class and helpers to interact with +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 -using MLFlow's tracking and model registry capabilities. - -The repository provides comprehensive functionality for: +Capabilities: - 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): - """ - Get the model URI based on the run_id. + """Build the artifact URI for a run's model. Args: - run_id (str): The run_id of the model. - prediction (bool): Whether to get prediction model URI (default: True) + run_id (str): MLflow run identifier. + 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: - """ - Get the run_id of a model based on its name and stage. + """Resolve the run_id for a registered model at a given stage. Args: - model_name (str): The name of the model. - stage (str): The stage of the model. + model_name (str): Registered model name. + 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 - artifact_path (str): Path to the artifact within the run + model_name (str): Registered model name. + 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. - 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. - If another type or format, must raise an error. + Normalize DataFrame index to the expected timestamp string format. + + The index must be timestamp-like. If the index is: + - 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 - metadata (dict): Metadata for logging + data (pd.DataFrame): DataFrame with timestamp index. + 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 - retention (int): Retention time in minutes + cache (dict): Cached model data with a 'timestamp' key. + 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_type (str): Type of model ('predict' or 'transform') - compressed (bool): Whether model is compressed - retention_target (str): Retention target ('model' or 'artifact') - cache (dict): Cached model data + model_name (str): Name of the model (for logging/consistency). + cache (dict): Cached model data structure. 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_key (str): Cache key for the model + model_name (str): Name of 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 - data (pd.DataFrame): Data to transform - retention (int): Cache retention time in minutes - flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') + model_name (str): Registered model name. + data (pd.DataFrame): Input data. + retention (int): Cache retention in minutes. + 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 - data (pd.DataFrame): Training data for model retraining - transform_flavor (str): Flavor for transformation model - predict_flavor (str): Flavor for prediction model - target_name (str): Target name - metadata (dict): Metadata for logging + model_name (str): Name of the MLflow model to retrain. + data (pd.DataFrame): Training data for model retraining. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + target_name (str | None): Optional target column; if None, use model target. + metadata (dict): Metadata for logging. Returns: - tuple: (prediction_model, data_model, experiment) - - prediction_model: Loaded prediction model for retraining - - data_model: Fitted transformation model - - experiment: MLFlow experiment name + dict[str, dict[str, Any]]: Mapping with prepared `prediction_model` and + `data_model`, including optional artifact paths. """ 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 - model_name (str): Name of the model being retrained - data (pd.DataFrame): Training data used for retraining - transform_flavor (str): Flavor for transformation model - predict_flavor (str): Flavor for prediction model - metadata (dict): Metadata for logging + experiment (str): MLflow experiment name for the retraining. + model_name (str): Name of the model being retrained. + data (pd.DataFrame): Training data used for retraining. + transform_flavor (str): Flavor for transformation model. + predict_flavor (str): Flavor for prediction model. + metadata (dict): Metadata for logging. Returns: - tuple: (status_message, experiment_name) - - status_message (str): Success confirmation message - - experiment_name (str): Name of the experiment + dict: Metadata about the created run and 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 - model_name (str): Name of the MLFlow model - metadata (dict): Metadata for logging + run_id (str): MLflow run ID containing the model to promote. + model_name (str): Name of the MLflow model. + 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 - in the MLFlow Model Registry in Production stage. + model_name (str): Name of the MLflow model to retrain. Must exist + 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: - - status_message (str): Success confirmation message or error details - - experiment_name (str): MLFlow experiment identifier for tracking + dict: Retraining operation results and experiment details. 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.