1ce8b9d3a7
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.
677 lines
26 KiB
Python
677 lines
26 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 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
|
|
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
|
|
|
|
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)
|
|
|
|
@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', {})
|
|
|
|
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='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)}',
|
|
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', {})
|
|
|
|
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='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)}',
|
|
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, log a new run, and register a model version.
|
|
|
|
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: Must include ``metadata``, ``model_name``, ``data`` (payload), and
|
|
``model_config`` with at least ``target``; optional ``full_retrain``, ``validation_fraction``.
|
|
|
|
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:
|
|
mv_src = self.mlflow_repository._client.get_model_version_by_alias(
|
|
name=model_name,
|
|
alias='production',
|
|
)
|
|
source_run_id = mv_src.run_id
|
|
|
|
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]:
|
|
"""
|
|
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)
|
|
|
|
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 {
|
|
'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:
|
|
mv = self.mlflow_repository._client.get_model_version_by_alias(
|
|
name=model_name,
|
|
alias='production',
|
|
)
|
|
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
|