aaf647efdf
Refactor MLFlow retraining logic to always use retrain method - Removed the conditional logic for full retraining, ensuring the `retrain` method is always called. - Updated the documentation in the `retrain_model` method to reflect the changes in retraining flow. - Adjusted tests to verify that the `retrain` method is invoked correctly, while ensuring `train` is not called when the full retrain flag is set.
688 lines
27 KiB
Python
688 lines
27 KiB
Python
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
|
|
import pandas as pd
|
|
from pandas import DataFrame, to_datetime
|
|
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
|
from sientia_do.notifications.models import NotificationLevel
|
|
from sientia_do.observability.logger import Logger
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.repository.minio_repository import MinioRepository
|
|
from sientia_do.temporal.constants import (
|
|
DATETIME_FORMAT,
|
|
DATETIME_FORMAT_MS_WITH_TZ,
|
|
DATETIME_FORMAT_WITH_TZ,
|
|
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
|
|
|
|
|
|
class MLFlow(MinioManager):
|
|
"""
|
|
Temporal activities that talk to MLflow through ``SientiaMLflowRepository`` and ``SientiaModel`` wrappers.
|
|
|
|
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``.
|
|
|
|
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_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
|
|
_DEFAULT_MODEL_ALIAS = 'production'
|
|
|
|
def __init__(
|
|
self,
|
|
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,
|
|
):
|
|
"""
|
|
Attach shared MLflow and MinIO clients used by all ML activities in this mixin.
|
|
|
|
Args:
|
|
- 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.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
|
|
MinioManager.__init__(
|
|
self, minio_repository, logger, notification_handler, metrics_controller
|
|
)
|
|
self.mlflow_repository = mlflow_repository
|
|
self.plugin_store = plugin_store
|
|
|
|
def close(self) -> None:
|
|
"""
|
|
Release MinIO manager resources held by the mixin.
|
|
|
|
Return:
|
|
None
|
|
"""
|
|
MinioManager.close(self)
|
|
|
|
def __del__(self):
|
|
self.close()
|
|
|
|
def _debug_dataframe(self, message: str, data: Any, metadata: dict[str, Any]) -> None:
|
|
"""
|
|
Log dataframe content only when row count is below the configured threshold
|
|
|
|
Args:
|
|
- message (str): Base log message to identify the dataframe in logs
|
|
- data (Any): Dataframe-like object expected to expose shape and to_csv
|
|
- metadata (dict[str, Any]): Workflow metadata for contextual logging
|
|
"""
|
|
self.debug(
|
|
build_dataframe_debug_message(
|
|
message=message,
|
|
data=data,
|
|
max_rows=self._MAX_DEBUG_DATAFRAME_ROWS,
|
|
),
|
|
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)
|
|
|
|
def _resolve_model_alias(self, model_config: dict[str, Any] | None = None) -> str:
|
|
"""
|
|
Resolve which MLflow alias should be used for model lookup/promotion.
|
|
|
|
Args:
|
|
- model_config: Optional model configuration that may include ``alias``.
|
|
|
|
Return:
|
|
str: Alias name trimmed and normalized; defaults to ``production``.
|
|
"""
|
|
if not model_config:
|
|
return self._DEFAULT_MODEL_ALIAS
|
|
alias = str(model_config.get('alias', self._DEFAULT_MODEL_ALIAS)).strip()
|
|
return alias or self._DEFAULT_MODEL_ALIAS
|
|
|
|
@activity.defn(name='request_transform')
|
|
async def request_transform(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
|
"""
|
|
Pivot long-format sensor rows, load the production wrapper, and run ``wrapper.transform``.
|
|
|
|
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: Dict with ``metadata``, ``model_name``, ``data`` (``MinioDataFramePayload``
|
|
dict or inline dataframe dict), and optional ``model_config``.
|
|
|
|
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)
|
|
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
model_alias = self._resolve_model_alias(model_config)
|
|
|
|
self._debug_dataframe('Raw input data:', data, metadata)
|
|
|
|
# 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'
|
|
)
|
|
|
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
|
data.fillna(np.nan, inplace=True)
|
|
|
|
data.columns.name = None
|
|
data.index.name = None
|
|
|
|
data['timestamp'] = data.index
|
|
|
|
self._debug_dataframe('Processed input data:', data, metadata)
|
|
|
|
try:
|
|
wrapper = self.mlflow_repository.get_cached_model(
|
|
model_name=model_name,
|
|
alias=model_alias,
|
|
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)}',
|
|
metadata,
|
|
)
|
|
|
|
self.debug(
|
|
f'Transform response data: \n {create_sample_dict(response_data, max_items=5, max_depth=5)}',
|
|
metadata,
|
|
)
|
|
|
|
self.info('Data transformed successfully', metadata)
|
|
|
|
if not response_data.get('success', False):
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=None,
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='transform',
|
|
status=response_data,
|
|
workflow_metadata=metadata,
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=self.logger,
|
|
)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=response_data['content'],
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='transform',
|
|
workflow_metadata=metadata,
|
|
status={
|
|
'success': True,
|
|
},
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=self.logger,
|
|
)
|
|
|
|
@activity.defn(name='request_predict')
|
|
async def request_predict(self, input_data: dict[str, Any]) -> MinioDataFramePayload:
|
|
"""
|
|
Load the production wrapper and call ``wrapper.predict`` on the prepared feature frame.
|
|
|
|
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: Same envelope as ``request_transform`` (``metadata``, ``model_name``,
|
|
``data``, optional ``model_config`` with ``retention_minutes``).
|
|
|
|
Return:
|
|
``MinioDataFramePayload`` with predictions or error status mirroring transform behaviour.
|
|
"""
|
|
metadata = input_data['metadata']
|
|
self.info('Predicting data...', metadata)
|
|
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
model_alias = self._resolve_model_alias(model_config)
|
|
|
|
self._debug_dataframe('Input data for prediction:', data, metadata)
|
|
|
|
input_index = data.index
|
|
|
|
data.replace(np.nan, None, inplace=True)
|
|
|
|
data['timestamp'] = data.index
|
|
data['timestamp'] = to_datetime(
|
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
|
).dt.strftime(DATETIME_FORMAT)
|
|
|
|
try:
|
|
wrapper = self.mlflow_repository.get_cached_model(
|
|
model_name=model_name,
|
|
alias=model_alias,
|
|
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)}',
|
|
metadata,
|
|
)
|
|
|
|
self.info('Data predicted successfully', metadata)
|
|
|
|
if not response_data.get('success', False):
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=None,
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='predict',
|
|
status=response_data,
|
|
workflow_metadata=metadata,
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=self.logger,
|
|
)
|
|
|
|
return await MinioDataFramePayload.from_dataframe(
|
|
dataframe=response_data['content'],
|
|
minio_repo=self.minio_repository,
|
|
model_name=model_name,
|
|
operation='predict',
|
|
workflow_metadata=metadata,
|
|
status={
|
|
'success': True,
|
|
},
|
|
last_timestamp=payload.last_timestamp,
|
|
logger=self.logger,
|
|
)
|
|
|
|
@activity.defn(name='retrain_model')
|
|
async def retrain_model(self, input_data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
Fit an updated wrapper from historical data, then log and register in MLflow.
|
|
|
|
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 →
|
|
run ``wrapper.retrain`` outside run timing → ``start_run`` with retrain tags → log input
|
|
CSV artifact → ``store_model`` and ``log_params``. Does not promote; the workflow calls
|
|
``update_production_model`` after validation.
|
|
|
|
Args:
|
|
- input_data: Must include ``metadata``, ``model_name``, ``data`` (payload), and
|
|
``model_config`` with at least ``target``.
|
|
|
|
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:
|
|
raise ValueError('Minio repository not initialized')
|
|
|
|
metadata = input_data['metadata']
|
|
|
|
try:
|
|
payload = MinioDataFramePayload.from_dict(input_data['data'])
|
|
data = await payload.retrieve(self.minio_repository, metadata)
|
|
|
|
except Exception as e:
|
|
trace = traceback.format_exc()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='ERROR_LOADING_RETRAIN_DATA',
|
|
message=f'Error loading retrain data: {e}',
|
|
block='retrain_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata)
|
|
return {
|
|
'success': False,
|
|
'message': f'Error loading retrain data: {e}',
|
|
'traceback': trace,
|
|
'timestamp': now().strftime(DATETIME_FORMAT_MS_WITH_TZ),
|
|
}
|
|
|
|
self.debug(f'Retrain data loaded successfully: shape {data.shape}', metadata)
|
|
|
|
model_name = input_data['model_name']
|
|
model_config = input_data.get('model_config', {})
|
|
|
|
self.info(f'Retraining model {model_name}...', metadata)
|
|
|
|
timestamp = data['timestamp'].max()
|
|
self.debug(f'Timestamp: {timestamp}', metadata)
|
|
|
|
if 'created_at' in data.columns:
|
|
data = data.sort_values('created_at', ascending=False).drop_duplicates(
|
|
subset=['variable', 'timestamp'], keep='first'
|
|
)
|
|
else:
|
|
data = data.drop_duplicates(subset=['variable', 'timestamp'], keep='first')
|
|
|
|
data.drop(columns=['model_id'], inplace=True, errors='ignore')
|
|
data.drop(columns=['created_at'], inplace=True, errors='ignore')
|
|
|
|
data = data.pivot(index='timestamp', columns='variable', values='value')
|
|
data.fillna(np.nan, inplace=True)
|
|
data.columns.name = None
|
|
|
|
data['timestamp'] = data.index
|
|
data['timestamp'] = to_datetime(
|
|
data['timestamp'], format=DATETIME_FORMAT_WITH_TZ
|
|
).dt.strftime(DATETIME_FORMAT)
|
|
data['timestamp'] = to_datetime(data['timestamp'], format=DATETIME_FORMAT)
|
|
|
|
data.columns.name = None
|
|
|
|
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),
|
|
}
|
|
|
|
try:
|
|
model_alias = self._resolve_model_alias(model_config)
|
|
mv_src = self.mlflow_repository._client.get_model_version_by_alias(
|
|
name=model_name,
|
|
alias=model_alias,
|
|
)
|
|
source_run_id = mv_src.run_id
|
|
|
|
wrapper = self.mlflow_repository.get_cached_model(
|
|
model_name=model_name,
|
|
alias=model_alias,
|
|
retention_minutes=0,
|
|
metadata=metadata,
|
|
)
|
|
|
|
# Keep heavy model fitting outside MLflow run timing.
|
|
wrapper.retrain(data)
|
|
|
|
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:
|
|
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]:
|
|
"""
|
|
Point the ``production`` alias at the model version registered for the retrain run.
|
|
|
|
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: ``metadata``, ``model_name``, and ``experiment`` with ``run_id`` and
|
|
``experiment_id`` (as returned from ``retrain_model``).
|
|
|
|
Return:
|
|
Dict with ``model_name``, promoted ``version``, ``mlflow_run_id``, ``mlflow_experiment_id``.
|
|
"""
|
|
metadata = input_data['metadata']
|
|
model_name = input_data['model_name']
|
|
experiment = input_data['experiment']
|
|
self.info(
|
|
f'Updating production model {model_name} from experiment {experiment}...', metadata
|
|
)
|
|
|
|
try:
|
|
run_id = experiment['run_id']
|
|
experiment_id = experiment['experiment_id']
|
|
|
|
version = self._resolve_model_version_for_run(run_id)
|
|
|
|
promote_alias = self._resolve_model_alias(input_data.get('model_config'))
|
|
self.mlflow_repository.promote_to_alias(
|
|
model_name=model_name,
|
|
version=version,
|
|
alias=promote_alias,
|
|
metadata=metadata,
|
|
)
|
|
|
|
self.info(f'Production model {model_name} updated successfully', metadata)
|
|
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()
|
|
await self.send_notification_async(
|
|
metadata=metadata,
|
|
notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
|
|
message=f'Error updating production model {model_name}: {e}',
|
|
block='update_production_model',
|
|
level=NotificationLevel.ERROR,
|
|
attachment_content=trace,
|
|
)
|
|
self.error(trace, metadata=metadata)
|
|
raise e
|
|
|
|
@activity.defn(name='get_reference_data')
|
|
async def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
|
|
"""
|
|
Download ``evaluation_data.csv`` from the MLflow run linked to ``production`` and parse it.
|
|
|
|
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: ``metadata`` and ``model_name`` for registry lookup.
|
|
|
|
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']
|
|
|
|
try:
|
|
model_alias = self._resolve_model_alias(input_data.get('model_config'))
|
|
mv = self.mlflow_repository._client.get_model_version_by_alias(
|
|
name=model_name,
|
|
alias=model_alias,
|
|
)
|
|
run_id = mv.run_id
|
|
|
|
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
|