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

from sientia_do.observability.metrics_controller import MetricsController
from temporalio import workflow

with workflow.unsafe.imports_passed_through():
    from typing import Any

    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import Logger
    from sientia_do.repository.minio_repository import MinioRepository
    from sientia_model.model_repository.mlflow_repository import SientiaMLflowRepository
    from sientia_model.model_repository.plugin_store import PluginStore

    from laborious.activities.api import API
    from laborious.activities.gates import Gates
    from laborious.activities.mlflow import MLFlow
    from laborious.activities.model_metrics import ModelMetrics
    from laborious.activities.opc import OPC
    from laborious.activities.storage import Storage
    from laborious.utils.connectors_config import build_mlflow_config


class Activities(Storage, MLFlow, Gates, OPC, ModelMetrics, API):
    """
    Central orchestrator for all Temporal activities used by Laborious workflows.

    Composes Storage (Postgres + MinIO offload), MLFlow (wrapper-based inference and retrain
    via ``SientiaMLflowRepository``), Gates (data quality and ML response filters), OPC exports,
    drift/simple metrics, and PI Web API writes. The worker constructs one ``Activities`` instance
    per process and registers its callables on multiple workers bound to different task queues.

    MLflow connectivity: unless ``mlflow_repository`` is injected (tests only), this class builds
    ``SientiaMLflowRepository`` from ``build_mlflow_config()`` so tracking credentials and URL
    stay aligned with the rest of Laborious env-based configuration.

    Attributes:
        Inherits and exposes behaviour from mixins; the MLFlow mixin holds ``mlflow_repository``
        and ``plugin_store`` after ``__init__``.
    """

    def __init__(
        self,
        postgres_config: dict[str, Any],
        plugin_store: PluginStore,
        minio_config: dict[str, Any],
        opc_config: dict[str, Any],
        pi_web_api_config: dict[str, Any],
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController | None = None,
        mlflow_repository: SientiaMLflowRepository | None = None,
    ):
        """
        Wire Postgres, MinIO, MLflow, OPC, gates, metrics, and PI Web API into a single object.

        A single ``MetricsController`` instance is created (or reused) and passed to MinIO,
        MLflow repository, and all mixins so Prometheus and SDK metrics stay consistent.

        Args:
         - postgres_config: Host, port, credentials, db name, and pool bounds for Storage.
         - plugin_store: ``PluginStore`` instance; the worker must call ``install_runtime`` before
           activities run so wrapper code is importable.
         - minio_config: Endpoint, keys, bucket, retention, and TLS flag for object storage payloads.
         - opc_config: Map of OPC server id to connection settings for ``OPC`` mixin.
         - pi_web_api_config: Base URL and auth for ``API`` mixin.
         - logger: Structured logger used across all activities.
         - notification_handler: Handler for alerts and persisted notifications.
         - metrics_controller: Optional shared controller; if ``None``, a new one is created.
         - mlflow_repository: Optional ``SientiaMLflowRepository`` for unit/e2e tests; in production
           leave unset so the repository is built from environment via ``build_mlflow_config()``.

        Raises:
            Exception: If any parent ``__init__`` fails (e.g. invalid config keys).

        Return:
            None
        """

        mc = metrics_controller or MetricsController(logger=logger)

        # Production path: one shared MLflow client for all model registry / tracking calls.
        if mlflow_repository is None:
            mlflow_cfg = build_mlflow_config()
            mlflow_repository = SientiaMLflowRepository(
                host=mlflow_cfg['url'],
                username=mlflow_cfg['username'],
                password=mlflow_cfg['password'],
                logger=logger,
                notification_handler=notification_handler,
                metrics_controller=mc,
            )

        minio_repository = MinioRepository(
            endpoint=minio_config['endpoint_url'],
            access_key=minio_config['access_key'],
            secret_key=minio_config['secret_key'],
            bucket=minio_config['default_bucket'],
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
            secure=minio_config['secure'],
        )

        Storage.__init__(
            self,
            host=postgres_config['host'],
            port=postgres_config['port'],
            user=postgres_config['user'],
            password=postgres_config['password'],
            dbname=postgres_config['dbname'],
            min_connections=postgres_config['min_connections'],
            max_connections=postgres_config['max_connections'],
            retention_hours=minio_config['retention_hours'],
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

        MLFlow.__init__(
            self,
            mlflow_repository=mlflow_repository,
            plugin_store=plugin_store,
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

        Gates.__init__(
            self,
            minio_repository=minio_repository,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

        OPC.__init__(
            self,
            opc_servers=opc_config,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

        ModelMetrics.__init__(
            self,
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

        API.__init__(
            self,
            base_url=pi_web_api_config['base_url'],
            auth_type=pi_web_api_config['auth_type'],
            auth_token=pi_web_api_config['auth_token'],
            logger=logger,
            notification_handler=notification_handler,
            metrics_controller=mc,
        )

    async def shutdown(self):
        """
        Close database pools, sync clients, and OPC sessions in a defined order.

        Should be invoked on worker exit so connection pools and OPC sessions are released
        cleanly before process termination.

        Return:
            None
        """
        Storage.close(self)
        MLFlow.close(self)
        Gates.close(self)
        await OPC.close(self)
        ModelMetrics.close(self)
        API.close(self)