sientia-dataops-laborious_temporal/laborious/activities/mlflow.py

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.observability.sientia_monitoring import SientiaMonitoring
    from sientia_do.repository.minio_repository_sync 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


class MLFlow(SientiaMonitoring):
    """
    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'
    _REFERENCE_ARTIFACT_CANDIDATES = ('retrain_input.csv', 'test_data.csv')

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

        self.minio_repository = minio_repository
        SientiaMonitoring.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=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
        """
        if self.minio_repository is not None:
            try:
                self.minio_repository.close()
            finally:
                self.minio_repository = None
        SientiaMonitoring.shutdown(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')
    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 = 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 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 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')
    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, sets the row index
        the same way as ``retrain_model`` (UTC ``DatetimeIndex`` from ``DATETIME_FORMAT_WITH_TZ``),
        restores that index on the prediction frame, normalizes the prediction index to
        ``DATETIME_FORMAT_WITH_TZ`` strings like ``request_transform``, and records ``response_time``.
        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 = 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)

        data.replace(np.nan, None, inplace=True)

        data.index = pd.DatetimeIndex(
            to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
        )
        input_index = data.index

        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()
            predict_data = self._detect_and_parse_datetime_index(predict_data, metadata)

            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 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 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')
    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 = payload.retrieve(self.minio_repository, metadata)

        except Exception as e:
            trace = traceback.format_exc()
            self.send_notification(
                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.index.name = None

        data.index = pd.DatetimeIndex(
            to_datetime(data.index, format=DATETIME_FORMAT_WITH_TZ, utc=True)
        )

        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')
    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()
            self.send_notification(
                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

    def _resolve_reference_artifact_name(self, run_id: str) -> str | None:
        """
        Pick the first available reference CSV artifact path from the MLflow run.

        Candidates are checked in priority order: ``retrain_input.csv``, then ``train_data.csv``.
        A path matches when it equals the candidate or ends with ``/<candidate>`` for nested layouts.

        Args:
         - run_id: MLflow run UUID linked to the production model version.

        Return:
            Artifact path string for ``download_artifacts``, or ``None`` if no candidate exists.
        """
        listed = self.mlflow_repository._client.list_artifacts(run_id)
        paths = [file_info.path for file_info in listed]
        for candidate in self._REFERENCE_ARTIFACT_CANDIDATES:
            for path in paths:
                if path == candidate or path.endswith(f'/{candidate}'):
                    return path
        return None

    def _find_downloaded_csv(self, tmpdir: str, artifact_name: str) -> Path | None:
        """
        Locate a downloaded reference CSV in the temp directory.

        Args:
         - tmpdir: Directory where ``download_artifacts`` wrote files.
         - artifact_name: Basename of the resolved artifact (e.g. ``retrain_input.csv``).

        Return:
            ``Path`` to the CSV file if found, else ``None``.
        """
        direct = Path(tmpdir) / artifact_name
        if direct.exists():
            return direct
        matches = list(Path(tmpdir).rglob(artifact_name))
        return matches[0] if matches else None

    @activity.defn(name='get_reference_data')
    def get_reference_data(self, input_data: dict[str, Any]) -> list[dict] | None:
        """
        Download reference training CSV from the MLflow run linked to the production alias.

        Resolves ``retrain_input.csv`` or ``train_data.csv`` via artifact listing before download.
        ``retrain_input.csv`` is preferred when both exist (most recent retrain snapshot). Used by
        drift workflows to compare live data against the reference distribution logged with the model.
        Timestamps are normalized to ``DATETIME_FORMAT`` string columns before returning records.

        Args:
         - input_data: ``metadata``, ``model_name``, and optional ``model_config`` with ``alias``.

        Return:
            List of row dicts with normalized timestamps, or ``None`` if resolution or load 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

            artifact_path = self._resolve_reference_artifact_name(run_id)
            if artifact_path is None:
                self.warning(f'Reference data not found for model {model_name}', metadata)
                return None

            artifact_name = Path(artifact_path).name
            tmpdir = tempfile.mkdtemp(prefix='laborious_eval_')
            try:
                self.mlflow_repository.download_artifacts(
                    run_id=run_id,
                    artifact_path=artifact_path,
                    dst_path=tmpdir,
                    metadata=metadata,
                )
                csv_path = self._find_downloaded_csv(tmpdir, artifact_name)
                if csv_path is None:
                    self.warning(f'Reference data not found for model {model_name}', metadata)
                    return None
                reference_data = pd.read_csv(csv_path)
            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