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.
This commit is contained in:
vitor-aignosi
2025-10-16 11:01:36 -03:00
parent de47820c4a
commit 644a43093a
3 changed files with 291 additions and 190 deletions

201
README.md
View File

@@ -1,42 +1,108 @@
# Sientia DataOps Laborious # 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 ## Features
### Core Functionality ### 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 ### 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 - **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 ## 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 ### Architecture Principles
#### 1. **Separation of Concerns** #### 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** #### 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** #### 3. **Scalability & Performance**
- **Horizontal Scaling**: Multiple worker instances for load distribution - **Horizontal Scaling**: Multiple worker instances for load distribution
@@ -53,66 +119,34 @@ The Laborious system uses a Temporal-based workflow architecture with clear sepa
### Key Components ### Key Components
#### **Worker (`laborious/worker/worker.py`)** #### **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/`)** #### **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 (`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/`)** #### **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 ### Data Flow Architecture
#### **1. Batch Prediction Pipeline** #### **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** #### **2. Model Retraining Pipeline**
@@ -121,33 +155,21 @@ Training Data → Model Retraining → Quality Validation →
Production Update → Notification & Monitoring Production Update → Notification & Monitoring
``` ```
#### **3. Real-time Export Pipeline**
```
Prediction Results → Data Formatting → OPC Server Write →
Success/Failure Metrics → Notification System
```
### Security Architecture ### Security Architecture
#### **Authentication & Authorization** #### **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** #### **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 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`) ### 1. Predictions Batch Workflow (`predictions_batch.py`)
@@ -351,8 +373,9 @@ flowchart LR
- Temporal server/cluster - Temporal server/cluster
- PostgreSQL database - PostgreSQL database
- MLFlow server - MLFlow server
- OPC server(s) - MinIO object storage (for MLFlow artifacts)
- MongoDB server (for notifications) - MongoDB server (for notifications)
- OPC server(s) if using OPC export
**Note**: External dependencies must be available either through: **Note**: External dependencies must be available either through:
- Kubernetes cluster deployment - Kubernetes cluster deployment

View File

@@ -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 io import BytesIO
from typing import Any from typing import Any
@@ -10,6 +18,21 @@ from sientia_do.observability.logger import Logger
class MinioRepository: 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__( def __init__(
self, self,
minio_endpoint_url: str, minio_endpoint_url: str,
@@ -20,6 +43,17 @@ class MinioRepository:
logger: Logger, logger: Logger,
notification_handler: NotificationHandler, 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 # MinIO settings shared with pandas s3fs
self.storage_options = { self.storage_options = {
'key': minio_access_key, 'key': minio_access_key,
@@ -54,11 +88,14 @@ class MinioRepository:
self.notification_handler = notification_handler self.notification_handler = notification_handler
def close(self): def close(self):
"""Close the underlying S3 client."""
self.s3_client.close() self.s3_client.close()
def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None: 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: try:
@@ -71,6 +108,14 @@ class MinioRepository:
def store_dataframe_as_parquet( def store_dataframe_as_parquet(
self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any] 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.ensure_bucket_exists(metadata)
self.logger.custom_info(f'Storing dataframe as parquet in {uri}', 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) 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: 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) 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) response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)

View File

@@ -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 - Model loading and caching with retention policies
- Data transformation and prediction operations - Data transformation and prediction operations
- Model retraining workflows - Model retraining workflows
@@ -37,6 +37,15 @@ INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'pytorch'
def force_memory_release(logger: Logger): 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() gc.collect()
try: try:
@@ -48,6 +57,14 @@ def force_memory_release(logger: Logger):
class MLFlowRepository: class MLFlowRepository:
def __init__(self, host: str, username: str, password: str, logger: Logger): 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 # set tracking uri
mlflow.set_tracking_uri(host) mlflow.set_tracking_uri(host)
@@ -64,15 +81,15 @@ class MLFlowRepository:
""" """
def get_model_uri(self, run_id: str, prediction: bool = True): 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: 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: Returns:
str: The model URI. str: Artifact URI to the selected model within the run.
""" """
run_info = mlflow.get_run(run_id) run_info = mlflow.get_run(run_id)
if prediction: if prediction:
@@ -82,15 +99,14 @@ class MLFlowRepository:
return model_uri return model_uri
def get_model_run_id(self, model_name: str, stage: str = 'Production') -> str: 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: 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: 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 # Use search_registered_models instead of deprecated get_latest_versions
registered_models = self.client.search_registered_models( registered_models = self.client.search_registered_models(
@@ -120,14 +136,14 @@ class MLFlowRepository:
def get_next_run_name(self, model_name: str) -> str: 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 This method calculates the next sequential run number for a model
by searching existing runs and incrementing the count. It ensures by searching existing runs and incrementing the count. It ensures
unique run names for model training and retraining operations. unique run names for model training and retraining operations.
Args: Args:
model_name (str): The name of the MLFlow model model_name (str): The name of the MLflow model
Returns: Returns:
str: The next run name in format 'model_name-run_number' 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 self, experiment_name: str, create_if_not_exists: bool = False
) -> Experiment: ) -> 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 This method searches for an MLFlow experiment by name and
returns its unique identifier. It provides error handling returns its unique identifier. It provides error handling
for non-existent experiments. for non-existent experiments.
Args: Args:
experiment_name (str): Name of the MLFlow experiment experiment_name (str): Name of the MLflow experiment
Returns: Returns:
int: MLFlow experiment ID Experiment: MLflow experiment object
Raises: 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) experiment = mlflow.get_experiment_by_name(experiment_name)
@@ -166,7 +182,14 @@ class MLFlowRepository:
return experiment return experiment
def get_model_params(self, run_id: str): 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) run_info = mlflow.get_run(run_id)
return run_info.data.params return run_info.data.params
@@ -176,14 +199,14 @@ class MLFlowRepository:
def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str: 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: 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: 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') run_id = self.get_model_run_id(model_name=model_name, stage='Production')
output_dir = f'{ARTIFACTS_PATH}/{model_name}' 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: 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: Args:
model_name (str): The name of the model to download from the registry. 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. mlflow.pyfunc.PyFuncModel: The loaded predictive model.
Notes: 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. - Warnings during the model loading process are suppressed.
""" """
model_uri = f'models:/{model_name}/production' model_uri = f'models:/{model_name}/production'
@@ -230,7 +253,7 @@ class MLFlowRepository:
def load_transform_model(self, model_name: str, flavor: str) -> Any: 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 This method retrieves the latest production model run ID for the given
model name, constructs the model URI, and loads the model using MLflow. model name, constructs the model URI, and loads the model using MLflow.
@@ -241,7 +264,7 @@ class MLFlowRepository:
artifact_path (str | None): Path to compressed artifacts if model is compressed artifact_path (str | None): Path to compressed artifacts if model is compressed
Returns: Returns:
Any: The loaded model object, as returned by `mlflow.sklearn.load_model`. Any: The loaded model object, depending on the flavor used.
Raises: Raises:
Exception: If the model run ID or URI cannot be retrieved, or if the Exception: If the model run ID or URI cannot be retrieved, or if the
@@ -266,7 +289,7 @@ class MLFlowRepository:
self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False
) -> tuple[Any, str | None]: ) -> tuple[Any, str | None]:
""" """
Download model based on type (predict or transform). Download model based on type ("predict" or "transform").
Args: Args:
model_name (str): Name of the model to download model_name (str): Name of the model to download
@@ -275,7 +298,7 @@ class MLFlowRepository:
load_wrapper (bool): Whether to load wrapper load_wrapper (bool): Whether to load wrapper
Returns: 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( self.logger.info(
@@ -317,17 +340,19 @@ class MLFlowRepository:
def detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame: 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: 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: Returns:
pd.DataFrame: DataFrame with converted datetime index pd.DataFrame: DataFrame with converted datetime index.
""" """
index = data.index index = data.index
@@ -368,14 +393,14 @@ class MLFlowRepository:
def check_cache_retention(self, cache: dict, retention: int) -> bool: 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: 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: 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() current_time = datetime.now()
cache_time = cache['timestamp'] cache_time = cache['timestamp']
@@ -385,17 +410,14 @@ class MLFlowRepository:
def handle_valid_model(self, model_name: str, cache: dict) -> dict: 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: 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: 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') 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. Clean up outdated cached model and its artifacts.
Args: 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: Returns:
None None
@@ -419,16 +441,16 @@ class MLFlowRepository:
def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any: 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: Args:
model_name (str): Name of the model to retrieve 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') model_type (str): Type of model ('predict' or 'transform')
flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch') flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
Returns: Returns:
Any: Model object Any: Model object.
""" """
# Retention is 0, download a new model # Retention is 0, download a new model
if retention <= 0: if retention <= 0:
@@ -488,16 +510,16 @@ class MLFlowRepository:
self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str
) -> pd.DataFrame | ndarray: ) -> pd.DataFrame | ndarray:
""" """
Get transformed data using cached transform model. Execute a cached operation using the requested model.
Args: 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: Returns:
pd.DataFrame: Transformed data pd.DataFrame | ndarray: Operation result.
""" """
if operation not in ['transform', 'predict']: if operation not in ['transform', 'predict']:
raise ValueError("Invalid operation. Use 'transform' or 'predict'.") raise ValueError("Invalid operation. Use 'transform' or 'predict'.")
@@ -531,7 +553,7 @@ class MLFlowRepository:
target_name: str | None = None, target_name: str | None = None,
) -> dict[str, dict[str, Any]]: ) -> 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: This method sets up the complete environment for model retraining by:
1. Loading the current production prediction model 1. Loading the current production prediction model
@@ -541,18 +563,16 @@ class MLFlowRepository:
5. Setting up the MLFlow experiment context 5. Setting up the MLFlow experiment context
Args: 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: 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) self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata)
@@ -659,6 +679,14 @@ class MLFlowRepository:
return retrain_data return retrain_data
def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): 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'] model = model_data['model']
self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata) self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata)
@@ -688,7 +716,7 @@ class MLFlowRepository:
predict_flavor: str = 'sklearn', predict_flavor: str = 'sklearn',
) -> dict: ) -> 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: This method performs the actual model retraining by:
1. Starting a new MLFlow run with descriptive metadata 1. Starting a new MLFlow run with descriptive metadata
@@ -700,17 +728,15 @@ class MLFlowRepository:
Args: Args:
prediction_model: MLFlow prediction model to retrain prediction_model: MLFlow prediction model to retrain
data_model: MLFlow transformation 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: 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'] prediction_model = retrain_data['prediction_model']
@@ -795,16 +821,16 @@ class MLFlowRepository:
self, run_id: str, model_name: str, metadata: dict self, run_id: str, model_name: str, metadata: dict
) -> 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 This method promotes a model from a specific MLFlow run to
production stage. It handles model registration, versioning, production stage. It handles model registration, versioning,
and stage transitions with proper error handling. and stage transitions with proper error handling.
Args: 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: Returns:
dict: Model update metadata containing: dict: Model update metadata containing:
@@ -863,7 +889,7 @@ class MLFlowRepository:
5. Manages model lifecycle based on retention policy (cleanup artifacts if needed) 5. Manages model lifecycle based on retention policy (cleanup artifacts if needed)
Parameters: 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. data (pd.DataFrame): The input data to be transformed by the model.
model_config (dict): Model configuration parameters model_config (dict): Model configuration parameters
metadata (dict): Metadata for logging metadata (dict): Metadata for logging
@@ -934,7 +960,7 @@ class MLFlowRepository:
9. Handles any exceptions and returns structured error information 9. Handles any exceptions and returns structured error information
Parameters: 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. data (pd.DataFrame): The input data to make predictions on.
model_retention (int): Cache retention time in minutes (0 = no caching). model_retention (int): Cache retention time in minutes (0 = no caching).
model_config (dict): Model configuration parameters model_config (dict): Model configuration parameters
@@ -1031,15 +1057,13 @@ class MLFlowRepository:
data (pd.DataFrame): Training data for model retraining. Must contain data (pd.DataFrame): Training data for model retraining. Must contain
all features required by both transformation and all features required by both transformation and
prediction models, including target variable. 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 model_config (dict): Model configuration parameters
metadata (dict): Metadata for logging metadata (dict): Metadata for logging
Returns: 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: Raises:
mlflow.exceptions.MlflowException: If model not found in registry mlflow.exceptions.MlflowException: If model not found in registry
@@ -1129,7 +1153,7 @@ class MLFlowRepository:
3. Returns comprehensive update metadata 3. Returns comprehensive update metadata
Args: 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. Must be a valid experiment that exists in MLFlow.
model_name (str): Name of the MLFlow model to update. Must exist model_name (str): Name of the MLFlow model to update. Must exist
in the MLFlow Model Registry. in the MLFlow Model Registry.