gitea_admin/sientia-dataops-laborious_temporal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1231

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files

205
README.md
View File

@@ -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

58
laborious/utils/repository/minio_repository.py
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 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)

218
laborious/utils/repository/model_repository.py
View File

@@ -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.