SIENTIAPDE-1773
Enhance environment configuration and update dependencies - Added new environment variables for PluginStore and MLflow configuration in `.env.example`, including `RUNTIME`, `STORE_BASE_URL`, `STORE_OWNER`, `STORE_REPO`, `STORE_BRANCH`, `STORE_USERNAME`, `STORE_PASSWORD`, `STORE_CACHE_TTL_SECONDS`, `PYPI_SERVER`, `PYPI_USERNAME`, and `PYPI_PASSWORD`. - Updated `git-requirements-mapping.txt` to reflect changes in repository names. - Modified `requirements-light.txt` and `requirements.txt` to upgrade `sientia-dataops-library` to version 1.12.0 and `sientia-mlops-library` to version 0.8.1. - Updated `values.yaml` to include new environment variables for worker runtime and PluginStore configuration. - Refactored E2E tests to utilize new MLflow repository stubs and PluginStore mocks for improved testing accuracy.
This commit is contained in:
vitor-aignosi
@@ -1,11 +1,18 @@
|
||||
from temporalio import activity, workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import tempfile
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from shutil import rmtree
|
||||
from typing import Any
|
||||
|
||||
import mlflow
|
||||
import numpy as np
|
||||
from pandas import to_datetime
|
||||
import pandas as pd
|
||||
from pandas import DataFrame, to_datetime
|
||||
from sklearn.model_selection import train_test_split
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.notifications.models import NotificationLevel
|
||||
from sientia_do.observability.logger import Logger
|
||||
@@ -18,79 +25,71 @@ with workflow.unsafe.imports_passed_through():
|
||||
now,
|
||||
)
|
||||
from sientia_do.utils.formatters import create_sample_dict
|
||||
from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
|
||||
from sientia_model.model_repository.plugin_store import PluginStore
|
||||
|
||||
from laborious.utils.dataframe_debug import build_dataframe_debug_message
|
||||
from laborious.utils.models.minio_dataframe_payload import MinioDataFramePayload
|
||||
from laborious.utils.repository.minio_manager import MinioManager
|
||||
from laborious.utils.repository.model_repository import MLFlowRepository
|
||||
|
||||
|
||||
class MLFlow(MinioManager):
|
||||
"""
|
||||
MLFlow integration activities for model inference operations.
|
||||
Temporal activities that talk to MLflow through ``SientiaMLflowRepository`` and ``SientiaModel`` wrappers.
|
||||
|
||||
This class provides activities for interacting with MLFlow models, including
|
||||
data transformation and prediction operations. It handles authentication,
|
||||
data preprocessing, and model management with configurable retention policies.
|
||||
Models are resolved by registered name and the ``production`` alias (not by legacy stages or
|
||||
separate transform/predict flavors). ``get_cached_model`` loads or reuses a wrapper; inference
|
||||
uses ``wrapper.transform`` / ``wrapper.predict``; retrain uses ``wrapper.retrain`` or
|
||||
``wrapper.train`` plus ``store_model`` and registry promotion via ``promote_to_alias``.
|
||||
|
||||
The class implements comprehensive error handling and logging for all
|
||||
MLFlow operations, ensuring reliable model inference in production environments.
|
||||
Large inputs and outputs flow through ``MinioDataFramePayload`` when workflows offload parquet
|
||||
to MinIO. On failure, transform/predict still return a payload with ``success: False`` and
|
||||
error details for downstream gates.
|
||||
|
||||
Attributes:
|
||||
mlflow_host (str): MLFlow server hostname
|
||||
mlflow_port (int): MLFlow server port
|
||||
mlflow_username (str): MLFlow authentication username
|
||||
mlflow_password (str): MLFlow authentication password
|
||||
model_monitoring_repository (MLFlowRepository): Repository for MLFlow operations
|
||||
mlflow_repository: Client for tracking, registry, artifact download, and run lifecycle.
|
||||
plugin_store: Reference to the store (runtime is installed on the worker; reserved for
|
||||
future store-backed helpers).
|
||||
"""
|
||||
|
||||
_MAX_DEBUG_DATAFRAME_ROWS = 100
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mlflow_host: str,
|
||||
mlflow_port: int,
|
||||
mlflow_username: str,
|
||||
mlflow_password: str,
|
||||
mlflow_repository: SientiaMLflowRepository,
|
||||
plugin_store: PluginStore,
|
||||
minio_repository: MinioRepository | None = None,
|
||||
logger: Logger | None = None,
|
||||
notification_handler: NotificationHandler | None = None,
|
||||
metrics_controller: MetricsController | None = None,
|
||||
):
|
||||
"""
|
||||
Initialize MLFlow activities with server configuration.
|
||||
Attach shared MLflow and MinIO clients used by all ML activities in this mixin.
|
||||
|
||||
Args:
|
||||
mlflow_host: MLFlow server hostname or IP address
|
||||
mlflow_port: MLFlow server port number
|
||||
mlflow_username: Username for MLFlow authentication
|
||||
mlflow_password: Password for MLFlow authentication
|
||||
logger: Logger instance for observability and debugging
|
||||
notification_handler: Notification handler for alerts and monitoring
|
||||
- mlflow_repository: Repository built by ``Activities`` (or injected in tests).
|
||||
- plugin_store: Plugin store instance from worker bootstrap.
|
||||
- minio_repository: MinIO client for ``MinioDataFramePayload`` upload/download.
|
||||
- logger: Structured logger.
|
||||
- notification_handler: Notifications on hard failures where applicable.
|
||||
- metrics_controller: Shared metrics controller.
|
||||
|
||||
Raises:
|
||||
Exception: If MLFlowRepository initialization fails
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
|
||||
MinioManager.__init__(
|
||||
self, minio_repository, logger, notification_handler, metrics_controller
|
||||
)
|
||||
self.mlflow_host = mlflow_host
|
||||
self.mlflow_port = mlflow_port
|
||||
self.mlflow_username = mlflow_username
|
||||
self.mlflow_password = mlflow_password
|
||||
|
||||
self.model_monitoring_repository = MLFlowRepository(
|
||||
f'{mlflow_host}:{mlflow_port}',
|
||||
mlflow_username,
|
||||
mlflow_password,
|
||||
logger,
|
||||
notification_handler,
|
||||
metrics_controller,
|
||||
)
|
||||
self.mlflow_repository = mlflow_repository
|
||||
self.plugin_store = plugin_store
|
||||
|
||||
def close(self) -> None:
|
||||
"""
|
||||
Close the MLFlow activity and clean up resources.
|
||||
Release MinIO manager resources held by the mixin.
|
||||
|
||||
Return:
|
||||
None
|
||||
"""
|
||||
MinioManager.close(self)
|
||||
|
||||
@@ -115,35 +114,100 @@ class MLFlow(MinioManager):
|
||||
metadata,
|
||||
)
|
||||
|
||||
def _detect_and_parse_datetime_index(self, data: pd.DataFrame, metadata: dict) -> pd.DataFrame:
|
||||
"""
|
||||
Ensure the transform output index is homogeneous and encoded as ``DATETIME_FORMAT_WITH_TZ`` strings.
|
||||
|
||||
Accepts an all-string index (validated against the format), or all-``datetime`` /
|
||||
``Timestamp`` (naive timestamps are localized to UTC before formatting). Mixed element types
|
||||
or unsupported types raise ``ValueError`` with a message logged at info level.
|
||||
|
||||
Args:
|
||||
- data: DataFrame whose index carries the time dimension after transform.
|
||||
- metadata: Workflow metadata for log correlation.
|
||||
|
||||
Return:
|
||||
``pd.DataFrame``: Same frame with a normalized string index; empty frames are returned as-is.
|
||||
"""
|
||||
|
||||
if data.empty:
|
||||
self.info('Data is empty, skipping datetime index detection and parsing', metadata)
|
||||
return data
|
||||
|
||||
index = data.index
|
||||
index_type = type(index[0])
|
||||
|
||||
self.info(f'Index type: {index_type}', metadata)
|
||||
|
||||
message = (
|
||||
f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, '
|
||||
f'string in format {DATETIME_FORMAT_WITH_TZ}.'
|
||||
)
|
||||
|
||||
if not all(isinstance(i, index_type) for i in index):
|
||||
types = map(str, map(type, index))
|
||||
raise ValueError(f'{message}. Elements are {",".join(types)}')
|
||||
|
||||
if index_type is str:
|
||||
try:
|
||||
pd.to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ)
|
||||
except ValueError as e:
|
||||
raise ValueError(f'{message}. Unable to parse given date format: {e}') from e
|
||||
|
||||
elif index_type is datetime or index_type is pd.Timestamp:
|
||||
idx = data.index
|
||||
if hasattr(idx, 'tz') and idx.tz is None:
|
||||
data.index = idx.tz_localize('UTC') # type: ignore[attr-defined]
|
||||
|
||||
data.index = data.index.strftime(DATETIME_FORMAT_WITH_TZ) # type: ignore[attr-defined]
|
||||
else:
|
||||
raise ValueError(f'{message}. Got {index_type}.')
|
||||
|
||||
return data
|
||||
|
||||
def _resolve_model_version_for_run(self, run_id: str) -> str:
|
||||
"""
|
||||
Map an MLflow ``run_id`` to the latest registered model version that produced that run.
|
||||
|
||||
``search_model_versions`` may return multiple versions if the model was registered more than
|
||||
once for the same run; the highest numeric ``version`` wins so promotion targets the newest
|
||||
artifact set.
|
||||
|
||||
Args:
|
||||
- run_id: Run UUID from ``retrain_model`` / experiment payload.
|
||||
|
||||
Return:
|
||||
str: Registry version string acceptable by ``promote_to_alias``.
|
||||
|
||||
Raises:
|
||||
ValueError: If the filter returns no versions (model not registered for this run).
|
||||
"""
|
||||
|
||||
versions = self.mlflow_repository._client.search_model_versions(
|
||||
filter_string=f"run_id='{run_id}'"
|
||||
)
|
||||
if not versions:
|
||||
raise ValueError(f'No registered model version found for run_id={run_id}')
|
||||
latest = max(versions, key=lambda v: int(v.version))
|
||||
return str(latest.version)
|
||||
|
||||
@activity.defn(name='request_transform')
|
||||
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Transform input data using MLFlow models.
|
||||
Pivot long-format sensor rows, load the production wrapper, and run ``wrapper.transform``.
|
||||
|
||||
This activity processes input data through MLFlow model transformation,
|
||||
including data preprocessing, format conversion, and validation. It handles
|
||||
data deduplication, pivoting, and cleanup to ensure optimal model performance.
|
||||
|
||||
The transformation process includes:
|
||||
1. Data deduplication based on variable and timestamp
|
||||
2. Data pivoting for model input format
|
||||
3. Null value handling and cleanup
|
||||
4. MLFlow model transformation request
|
||||
5. Response validation and logging
|
||||
Expected tabular shape after load: columns including ``variable``, ``timestamp``, ``value``,
|
||||
and ``created_at`` for deduplication. Data are sorted by ``created_at``, de-duplicated per
|
||||
``(variable, timestamp)``, pivoted wide, then passed to the model. ``model_config`` may
|
||||
include ``retention_minutes`` for wrapper cache TTL.
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for transformation
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict): Input data for transformation
|
||||
- model_name (str): Name of the MLFlow model to use
|
||||
- model_retention (int): Model retention period in minutes
|
||||
- input_data: Dict with ``metadata``, ``model_name``, ``data`` (``MinioDataFramePayload``
|
||||
dict or inline dataframe dict), and optional ``model_config``.
|
||||
|
||||
Returns:
|
||||
dict: Transformed data from MLFlow model
|
||||
|
||||
Raises:
|
||||
Exception: If transformation fails or MLFlow model is unavailable
|
||||
Return:
|
||||
``MinioDataFramePayload`` with transformed frame and ``success: True``, or a payload
|
||||
with ``success: False`` and exception details in ``status`` if transform fails.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Transforming data...', metadata)
|
||||
@@ -156,12 +220,11 @@ class MLFlow(MinioManager):
|
||||
|
||||
self._debug_dataframe('Raw input data:', data, metadata)
|
||||
|
||||
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
||||
# Long → wide: keep newest row per (variable, timestamp), then pivot for the wrapper API.
|
||||
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
||||
subset=['variable', 'timestamp'], keep='first'
|
||||
)
|
||||
|
||||
# Pivot data for model input format
|
||||
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||
data.fillna(np.nan, inplace=True)
|
||||
|
||||
@@ -172,10 +235,25 @@ class MLFlow(MinioManager):
|
||||
|
||||
self._debug_dataframe('Processed input data:', data, metadata)
|
||||
|
||||
# Request transformation from MLFlow model
|
||||
response_data = await self.model_monitoring_repository.transform(
|
||||
model_name, data, model_config, metadata
|
||||
)
|
||||
try:
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias='production',
|
||||
retention_minutes=model_config.get('retention_minutes', 0),
|
||||
metadata=metadata,
|
||||
)
|
||||
transformed_df, transform_meta = wrapper.transform(data)
|
||||
if transform_meta:
|
||||
self.info(f'Wrapper transform metadata: {transform_meta}', metadata)
|
||||
|
||||
transformed_df = self._detect_and_parse_datetime_index(transformed_df, metadata)
|
||||
|
||||
response_data: dict[str, Any] = {'success': True, 'content': transformed_df}
|
||||
except Exception as e:
|
||||
response_data = {
|
||||
'success': False,
|
||||
'content': {'message': str(e), 'traceback': traceback.format_exc()},
|
||||
}
|
||||
|
||||
self.debug(
|
||||
f'Transform raw response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||
@@ -217,32 +295,19 @@ class MLFlow(MinioManager):
|
||||
@activity.defn(name='request_predict')
|
||||
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
||||
"""
|
||||
Execute predictions using MLFlow models.
|
||||
Load the production wrapper and call ``wrapper.predict`` on the prepared feature frame.
|
||||
|
||||
This activity performs ML model inference using MLFlow models with the
|
||||
transformed data. It handles data format conversion, null value processing,
|
||||
and model prediction requests with comprehensive error handling.
|
||||
|
||||
The prediction process includes:
|
||||
1. Data format validation and cleanup
|
||||
2. Null value handling for model compatibility
|
||||
3. MLFlow model prediction request
|
||||
4. Response validation and logging
|
||||
5. Performance monitoring and metrics
|
||||
The activity normalizes ``NaN`` to ``None`` for JSON-friendly columns, rebuilds a
|
||||
``timestamp`` column in the internal string format, preserves the original index for
|
||||
alignment, and records ``response_time`` seconds on the output frame. Non-DataFrame
|
||||
predictions are coerced to a single ``prediction`` column.
|
||||
|
||||
Args:
|
||||
input_data: Configuration and data for prediction
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict): Transformed data for prediction
|
||||
- model_name (str): Name of the MLFlow model to use
|
||||
- model_retention (int): Model retention period in minutes
|
||||
- input_data: Same envelope as ``request_transform`` (``metadata``, ``model_name``,
|
||||
``data``, optional ``model_config`` with ``retention_minutes``).
|
||||
|
||||
Returns:
|
||||
dict: Prediction results from MLFlow model
|
||||
|
||||
Raises:
|
||||
Exception: If prediction fails or MLFlow model is unavailable
|
||||
Return:
|
||||
``MinioDataFramePayload`` with predictions or error status mirroring transform behaviour.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
self.info('Predicting data...', metadata)
|
||||
@@ -255,7 +320,8 @@ class MLFlow(MinioManager):
|
||||
|
||||
self._debug_dataframe('Input data for prediction:', data, metadata)
|
||||
|
||||
# Convert numpy.nan to None for model compatibility
|
||||
input_index = data.index
|
||||
|
||||
data.replace(np.nan, None, inplace=True)
|
||||
|
||||
data['timestamp'] = data.index
|
||||
@@ -263,10 +329,41 @@ class MLFlow(MinioManager):
|
||||
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
||||
).dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
# Request prediction from MLFlow model
|
||||
response_data = await self.model_monitoring_repository.predict(
|
||||
model_name, data, model_config, metadata
|
||||
)
|
||||
try:
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias='production',
|
||||
retention_minutes=model_config.get('retention_minutes', 0),
|
||||
metadata=metadata,
|
||||
)
|
||||
start_time = datetime.now()
|
||||
predict_data, pred_meta = wrapper.predict({}, data)
|
||||
end_time = datetime.now()
|
||||
|
||||
if pred_meta:
|
||||
self.info(f'Wrapper predict metadata: {pred_meta}', metadata)
|
||||
|
||||
if isinstance(predict_data, DataFrame):
|
||||
self._debug_dataframe(
|
||||
'Data received from model prediction:', predict_data, metadata
|
||||
)
|
||||
predict_data.columns = pd.Index(['prediction'])
|
||||
else:
|
||||
self.debug(
|
||||
f'Data received from model prediction (not a DataFrame): {predict_data}',
|
||||
metadata,
|
||||
)
|
||||
predict_data = pd.DataFrame(predict_data, columns=['prediction'])
|
||||
|
||||
predict_data.index = input_index
|
||||
predict_data['response_time'] = (end_time - start_time).total_seconds()
|
||||
|
||||
response_data: dict[str, Any] = {'success': True, 'content': predict_data}
|
||||
except Exception as e:
|
||||
response_data = {
|
||||
'success': False,
|
||||
'content': {'message': str(e), 'traceback': traceback.format_exc()},
|
||||
}
|
||||
|
||||
self.debug(
|
||||
f'Prediction response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
||||
@@ -303,34 +400,22 @@ class MLFlow(MinioManager):
|
||||
@activity.defn(name='retrain_model')
|
||||
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Retrain MLFlow models with updated training data.
|
||||
Fit an updated wrapper from historical data, log a new run, and register a model version.
|
||||
|
||||
This activity orchestrates the complete model retraining process,
|
||||
including data preparation, model retraining execution, and result
|
||||
validation. It handles data preprocessing, column cleanup, and
|
||||
comprehensive error handling for production model management.
|
||||
|
||||
The retraining process includes:
|
||||
1. Data timestamp extraction and validation
|
||||
2. Column cleanup and data preparation
|
||||
3. Data pivoting for model input format
|
||||
4. MLFlow model retraining execution
|
||||
5. Result validation and error handling
|
||||
Flow: load long-format data from MinIO → dedupe/pivot like inference prep → require
|
||||
``model_config['target']`` → read current ``production`` version for ``source_run_id`` tag →
|
||||
``start_run`` with retrain tags → ``wrapper.retrain`` or ``wrapper.train`` when
|
||||
``full_retrain`` is set (optional ``validation_fraction``) → log input CSV artifact →
|
||||
``store_model`` and ``log_params``. Does not promote; the workflow calls
|
||||
``update_production_model`` after validation.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- data (dict[str, Any]): Training data for model retraining
|
||||
- model_name (str): Name of the MLFlow model to retrain
|
||||
- input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
|
||||
``model_config`` with at least ``target``; optional ``full_retrain``, ``validation_fraction``.
|
||||
|
||||
Returns:
|
||||
dict: Retraining results containing:
|
||||
- status (str): Retraining operation status
|
||||
- timestamp (str): Timestamp of the retraining operation
|
||||
- experiment (str): MLFlow experiment identifier
|
||||
|
||||
Raises:
|
||||
Exception: If retraining fails or encounters critical errors
|
||||
Return:
|
||||
On success: ``success``, ``experiment`` (``run_id``, ``experiment_id``, ``experiment_name``),
|
||||
``message``, ``timestamp``. On failure: ``success: False``, error fields, and optional trace.
|
||||
"""
|
||||
|
||||
if self.minio_repository is None:
|
||||
@@ -339,7 +424,6 @@ class MLFlow(MinioManager):
|
||||
metadata = input_data['metadata']
|
||||
|
||||
try:
|
||||
# Payload-based retrain input (inline dict or MinIO offloaded).
|
||||
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
||||
data = await payload.retrieve(self.minio_repository, metadata)
|
||||
|
||||
@@ -371,7 +455,6 @@ class MLFlow(MinioManager):
|
||||
timestamp = data['timestamp'].max()
|
||||
self.debug(f'Timestamp: {timestamp}', metadata)
|
||||
|
||||
# Sort by created_at in descending order and keep first occurrence of each variable/timestamp pair
|
||||
if 'created_at' in data.columns:
|
||||
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
||||
subset=['variable', 'timestamp'], keep='first'
|
||||
@@ -382,10 +465,8 @@ class MLFlow(MinioManager):
|
||||
data.drop(columns=['model_id'], inplace=True, errors='ignore')
|
||||
data.drop(columns=['created_at'], inplace=True, errors='ignore')
|
||||
|
||||
# Pivot data for model input format
|
||||
data = data.pivot(index='timestamp', columns='variable', values='value')
|
||||
data.fillna(np.nan, inplace=True)
|
||||
# data.reset_index(inplace=True)
|
||||
data.columns.name = None
|
||||
|
||||
data['timestamp'] = data.index
|
||||
@@ -396,60 +477,110 @@ class MLFlow(MinioManager):
|
||||
|
||||
data.columns.name = None
|
||||
|
||||
retrain_output = await self.model_monitoring_repository.retrain_model(
|
||||
data=data, model_name=model_name, model_config=model_config, metadata=metadata
|
||||
)
|
||||
target = model_config.get('target')
|
||||
if target is None:
|
||||
msg = 'model_config must include "target" for retraining'
|
||||
self.info(msg, metadata)
|
||||
return {
|
||||
'success': False,
|
||||
'experiment': None,
|
||||
'message': msg,
|
||||
'traceback': '',
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
if not retrain_output['success']:
|
||||
trace = retrain_output['traceback']
|
||||
await self.send_notification_async(
|
||||
metadata=metadata,
|
||||
notification_id='RETRAIN_MODEL_ERROR',
|
||||
message=f'Error retraining model {model_name}: {retrain_output["message"]}',
|
||||
block='retrain_model',
|
||||
level=NotificationLevel.ERROR,
|
||||
attachment_content=trace,
|
||||
try:
|
||||
mv_src = self.mlflow_repository._client.get_model_version_by_alias(
|
||||
name=model_name,
|
||||
alias='production',
|
||||
)
|
||||
self.error(trace, metadata=metadata)
|
||||
source_run_id = mv_src.run_id
|
||||
|
||||
return {**retrain_output, 'timestamp': timestamp}
|
||||
wrapper = self.mlflow_repository.get_cached_model(
|
||||
model_name=model_name,
|
||||
alias='production',
|
||||
retention_minutes=0,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
run_name = f'{model_name}-retrain-{datetime.now().strftime("%Y%m%d%H%M%S")}'
|
||||
|
||||
with self.mlflow_repository.start_run(
|
||||
model_name=model_name,
|
||||
run_name=run_name,
|
||||
experiment_name=model_name,
|
||||
tags={'retrain': 'true', 'source_run_id': source_run_id},
|
||||
metadata=metadata,
|
||||
) as run_info:
|
||||
if model_config.get('full_retrain'):
|
||||
val_frac = float(model_config.get('validation_fraction', 0.2))
|
||||
train_df, val_df = train_test_split(data, test_size=val_frac, random_state=42)
|
||||
wrapper.train(
|
||||
train_data=train_df,
|
||||
val_data=val_df,
|
||||
target=target,
|
||||
)
|
||||
else:
|
||||
wrapper.retrain(data)
|
||||
|
||||
tmp_dir = tempfile.mkdtemp(prefix='laborious_retrain_')
|
||||
try:
|
||||
raw_csv = Path(tmp_dir) / 'retrain_input.csv'
|
||||
data.to_csv(raw_csv, index=False)
|
||||
mlflow.log_artifact(str(raw_csv))
|
||||
finally:
|
||||
rmtree(tmp_dir, ignore_errors=True)
|
||||
|
||||
wrapper.store_model(name=model_name)
|
||||
|
||||
self.mlflow_repository.log_params(
|
||||
{
|
||||
'retrain': 'true',
|
||||
'retrain_date': datetime.now().isoformat(),
|
||||
'source_run_id': source_run_id,
|
||||
'retrain_samples': str(data.shape),
|
||||
}
|
||||
)
|
||||
|
||||
experiment_payload = {
|
||||
'run_id': run_info.run_id,
|
||||
'experiment_id': run_info.experiment_id,
|
||||
'experiment_name': model_name,
|
||||
}
|
||||
|
||||
return {
|
||||
'success': True,
|
||||
'experiment': experiment_payload,
|
||||
'message': 'Model retrained successfully.',
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
error_msg = f'Error retraining model {model_name}: {e}'
|
||||
self.info(error_msg, metadata)
|
||||
return {
|
||||
'success': False,
|
||||
'experiment': None,
|
||||
'message': error_msg,
|
||||
'traceback': traceback.format_exc(),
|
||||
'timestamp': str(timestamp),
|
||||
}
|
||||
|
||||
@activity.defn(name='update_production_model')
|
||||
async def update_production_model(self, input_data: dict[str, Any]) -> dict[Any, Any]:
|
||||
"""
|
||||
Update production model with newly trained model version.
|
||||
Point the ``production`` alias at the model version registered for the retrain run.
|
||||
|
||||
This activity manages the critical process of updating production
|
||||
models with newly trained versions. It handles model deployment,
|
||||
status tracking, and comprehensive reporting for operational
|
||||
visibility and audit trails.
|
||||
|
||||
The update process includes:
|
||||
1. Production model update execution
|
||||
2. Status and metadata tracking
|
||||
3. Comprehensive reporting and logging
|
||||
4. Error handling and notification
|
||||
5. Audit trail maintenance
|
||||
Resolves the highest numeric registry version whose ``run_id`` matches
|
||||
``experiment['run_id']``, then calls ``promote_to_alias``. On failure, sends a notification
|
||||
and re-raises so the workflow can surface the error.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_name (str): Name of the MLFlow model to update
|
||||
- experiment (str): MLFlow experiment identifier
|
||||
- model_id (str): Unique identifier for the model version
|
||||
- timestamp (str): Timestamp of the update operation
|
||||
- status (str): Current status of the model update
|
||||
- input_data: ``metadata``, ``model_name``, and ``experiment`` with ``run_id`` and
|
||||
``experiment_id`` (as returned from ``retrain_model``).
|
||||
|
||||
Returns:
|
||||
dict[Any, Any]: Comprehensive update report containing:
|
||||
- model_id (str): Model version identifier
|
||||
- model_name (str): Name of the updated model
|
||||
- timestamp (str): Update operation timestamp
|
||||
- status (str): Update operation status
|
||||
- Additional MLFlow response metadata
|
||||
|
||||
Raises:
|
||||
Exception: If production model update fails
|
||||
Return:
|
||||
Dict with ``model_name``, promoted ``version``, ``mlflow_run_id``, ``mlflow_experiment_id``.
|
||||
"""
|
||||
metadata = input_data['metadata']
|
||||
model_name = input_data['model_name']
|
||||
@@ -459,12 +590,25 @@ class MLFlow(MinioManager):
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self.model_monitoring_repository.update_production_model(
|
||||
experiment=experiment, model_name=model_name, metadata=metadata
|
||||
run_id = experiment['run_id']
|
||||
experiment_id = experiment['experiment_id']
|
||||
|
||||
version = self._resolve_model_version_for_run(run_id)
|
||||
|
||||
self.mlflow_repository.promote_to_alias(
|
||||
model_name=model_name,
|
||||
version=version,
|
||||
alias='production',
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
self.info(f'Production model {model_name} updated successfully', metadata)
|
||||
return response
|
||||
return {
|
||||
'model_name': model_name,
|
||||
'version': version,
|
||||
'mlflow_run_id': run_id,
|
||||
'mlflow_experiment_id': experiment_id,
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
trace = traceback.format_exc()
|
||||
@@ -482,47 +626,51 @@ class MLFlow(MinioManager):
|
||||
@activity.defn(name='get_reference_data')
|
||||
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
||||
"""
|
||||
Get reference data from the MLflow Model Registry.
|
||||
Download ``evaluation_data.csv`` from the MLflow run linked to ``production`` and parse it.
|
||||
|
||||
This method retrieves evaluation reference data stored as artifacts in the
|
||||
MLflow Model Registry. The reference data is typically used for model
|
||||
drift detection, performance comparison, and quality validation. The method
|
||||
loads the data from a CSV artifact file and formats timestamps for
|
||||
consistent processing.
|
||||
|
||||
The method handles:
|
||||
1. Loading evaluation data artifact from MLflow Model Registry
|
||||
2. Timestamp parsing and formatting for consistency
|
||||
3. Data conversion to dictionary format for workflow consumption
|
||||
4. Graceful handling of missing reference data
|
||||
Used by drift workflows to compare live data against the reference distribution logged with
|
||||
the model. Artifacts are downloaded to a temp directory, discovered via ``rglob`` (nested
|
||||
layout-safe), then timestamps are normalized to ``DATETIME_FORMAT`` string columns before
|
||||
returning record-oriented dicts.
|
||||
|
||||
Args:
|
||||
input_data (dict): Input data containing:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- model_name (str): Name of the MLFlow model to get reference data from
|
||||
- input_data: ``metadata`` and ``model_name`` for registry lookup.
|
||||
|
||||
Returns:
|
||||
list[dict[Hashable, Any]] | None: Reference data from the MLflow Model Registry
|
||||
as a list of dictionaries. Returns None if reference data is not found
|
||||
or if the artifact does not exist.
|
||||
|
||||
Raises:
|
||||
Exception: If artifact loading fails or encounters errors during processing
|
||||
Return:
|
||||
List of row dicts, or ``None`` if the artifact path is missing or any step fails.
|
||||
"""
|
||||
|
||||
metadata = input_data['metadata']
|
||||
model_name = input_data['model_name']
|
||||
artifact = 'evaluation_data.csv'
|
||||
|
||||
reference_data = await self.model_monitoring_repository.load_artifact_dataframe(
|
||||
model_name=model_name, artifact_path=artifact, metadata=metadata
|
||||
)
|
||||
try:
|
||||
mv = self.mlflow_repository._client.get_model_version_by_alias(
|
||||
name=model_name,
|
||||
alias='production',
|
||||
)
|
||||
run_id = mv.run_id
|
||||
|
||||
if reference_data is None:
|
||||
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||
tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
|
||||
try:
|
||||
self.mlflow_repository.download_artifacts(
|
||||
run_id=run_id,
|
||||
artifact_path='evaluation_data.csv',
|
||||
dst_path=tmpdir,
|
||||
metadata=metadata,
|
||||
)
|
||||
csv_candidates = list(Path(tmpdir).rglob('evaluation_data.csv'))
|
||||
if not csv_candidates:
|
||||
self.warning(f'Reference data not found for model {model_name}', metadata)
|
||||
return None
|
||||
|
||||
reference_data = pd.read_csv(csv_candidates[0])
|
||||
finally:
|
||||
rmtree(tmpdir, ignore_errors=True)
|
||||
|
||||
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
return reference_data.to_dict(orient='records')
|
||||
|
||||
except Exception as e:
|
||||
self.warning(f'Reference data not found for model {model_name}: {e}', metadata)
|
||||
return None
|
||||
|
||||
reference_data['timestamp'] = to_datetime(reference_data['timestamp'])
|
||||
reference_data['timestamp'] = reference_data['timestamp'].dt.strftime(DATETIME_FORMAT)
|
||||
|
||||
return reference_data.to_dict(orient='records')
|
||||
|
||||
Reference in New Issue
Block a user