SIENTIAPDE-1325

Refactor monitoring and metrics integration across various components

- Removed coverage options from `pyproject.toml`.
- Updated prediction metrics in `README.md` to replace `pipeline_name` with `workflow_name`.
- Upgraded `sientia-dataops-library` dependency version in `requirements-light.txt` and `requirements.txt`.
- Enhanced metrics handling in `laborious` activities, including `Activities`, `Gates`, `MLFlow`, and `OPC`, to utilize a new `MetricsController`.
- Refactored metric emission methods to improve clarity and consistency across the codebase.
- Updated tests to reflect changes in metrics handling and ensure proper functionality.

laborious/utils/repository/minio_repository.py

@@ -6,6 +6,7 @@ object storage using boto3. It supports creating buckets on demand and
 storing/loading pandas DataFrames in Parquet format.
 """
 
+import time
 from io import BytesIO
 from typing import Any
 
@@ -15,9 +16,13 @@ from botocore.exceptions import ClientError
 from pandas import DataFrame, read_parquet
 from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
 from sientia_do.observability.logger import Logger
+from sientia_do.observability.metrics_controller import MetricsController
+from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+from laborious import metrics
 
 
-class MinioRepository:
+class MinioRepository(SientiaMonitoring):
     """
     Repository for interacting with a MinIO (S3-compatible) object storage.
 
@@ -43,6 +48,7 @@ class MinioRepository:
         minio_default_bucket: str,
         logger: Logger,
         notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
     ):
         """Initialize the repository and S3 client.
 
@@ -55,6 +61,7 @@ class MinioRepository:
             logger (Logger): Logger instance for structured logs.
             notification_handler (NotificationHandler): Notification handler.
         """
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
         # MinIO settings shared with pandas s3fs
         self.storage_options = {
             'key': minio_access_key,
@@ -85,28 +92,58 @@ class MinioRepository:
             ),
         )
 
-        self.logger = logger
-        self.notification_handler = notification_handler
-
     def close(self):
         """Close the underlying S3 client."""
         self.s3_client.close()
 
-    def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None:
+    async def create_bucket(self, metadata: dict[str, Any]) -> None:
+        core_labels = {
+            **self.get_core_labels(metadata, operation_type='create_bucket'),
+            'bucket_name': self.minio_bucket,
+            'object_name': '-',
+        }
+        self.info(f"Creating bucket '{self.minio_bucket}'", metadata)
+
+        start_time = time.time()
+        try:
+            self.s3_client.create_bucket(Bucket=self.minio_bucket)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels)
+
+    async def ensure_bucket_exists(self, metadata: dict[str, Any]) -> None:
         """Ensure the default bucket exists; create it if missing.
 
         Args:
             metadata (dict[str, Any]): Metadata used for structured logging.
         """
+        self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata)
+        core_labels = {
+            **self.get_core_labels(metadata, operation_type='head_bucket'),
+            'bucket_name': self.minio_bucket,
+            'object_name': '-',
+        }
+        self.info(f"Checking if bucket '{self.minio_bucket}' exists", metadata)
+
+        start_time = time.time()
 
         try:
-            self.logger.custom_info(f"Checking if bucket '{self.minio_bucket}' exists", metadata)
             self.s3_client.head_bucket(Bucket=self.minio_bucket)
         except ClientError:
-            self.logger.custom_info(f"Creating bucket '{self.minio_bucket}'", metadata)
-            self.s3_client.create_bucket(Bucket=self.minio_bucket)
+            await self.create_bucket(metadata)
 
-    def store_dataframe_as_parquet(
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        else:
+            await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels)
+            await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels)
+
+    async def store_dataframe_as_parquet(
         self, dataframe: DataFrame, uri: str, object_name: str, metadata: dict[str, Any]
     ):
         """Persist a DataFrame as a Parquet object in the default bucket.
@@ -117,18 +154,36 @@ class MinioRepository:
             object_name (str): Object key (path/key within the bucket).
             metadata (dict[str, Any]): Metadata used for structured logging.
         """
-        self.ensure_bucket_exists(metadata)
+        await self.ensure_bucket_exists(metadata)
 
-        self.logger.custom_info(f'Storing dataframe as parquet in {uri}', metadata)
+        self.info(f'Storing dataframe as parquet in {uri}', metadata)
 
         buffer = BytesIO()
         dataframe.to_parquet(buffer, engine='pyarrow', index=True)
         buffer.seek(0)
-        self.s3_client.put_object(Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue())
 
-        self.logger.custom_info(f'Dataframe stored as parquet in {uri}', metadata)
+        core_labels = {
+            **self.get_core_labels(metadata, operation_type='put_object'),
+            'bucket_name': self.minio_bucket,
+            'object_name': object_name,
+        }
+        start_time = time.time()
+        try:
+            self.s3_client.put_object(
+                Bucket=self.minio_bucket, Key=object_name, Body=buffer.getvalue()
+            )
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
 
-    def get_parquet_as_dataframe(self, object_key: str, metadata: dict[str, Any]) -> DataFrame:
+        await self.observe_lag(start_time, metrics.MINIO_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MINIO_WRITE_COUNT, tags=core_labels)
+
+        self.info(f'Dataframe stored as parquet in {uri}', metadata)
+
+    async def get_parquet_as_dataframe(
+        self, object_key: str, metadata: dict[str, Any]
+    ) -> DataFrame:
         """Load a Parquet object from the default bucket into a DataFrame.
 
         Args:
@@ -138,9 +193,22 @@ class MinioRepository:
         Returns:
             DataFrame: Loaded DataFrame.
         """
-        self.logger.custom_info(f'Getting parquet as dataframe from {object_key}', metadata)
+        self.info(f'Getting parquet as dataframe from {object_key}', metadata)
 
-        response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)
+        core_labels = {
+            **self.get_core_labels(metadata, operation_type='get_object'),
+            'bucket_name': self.minio_bucket,
+            'object_name': object_key,
+        }
+        start_time = time.time()
+        try:
+            response = self.s3_client.get_object(Bucket=self.minio_bucket, Key=object_key)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MINIO_READ_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MINIO_READ_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MINIO_READ_COUNT, tags=core_labels)
 
         # Read the content into a BytesIO buffer to support seek operations
         buffer = BytesIO(response['Body'].read())

laborious/utils/repository/model_repository.py

@@ -17,6 +17,7 @@ Capabilities:
 import ctypes
 import gc
 import threading
+import time
 import traceback
 from datetime import datetime, timedelta
 from os import environ, makedirs, path
@@ -27,9 +28,14 @@ import mlflow
 import pandas as pd
 from mlflow.entities import Experiment
 from numpy import ndarray
+from sientia_do.notifications.handlers import NotificationHandler
 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.temporal.constants import DATETIME_FORMAT_WITH_TZ
 
+from laborious import metrics
+
 ARTIFACTS_PATH = './tmp/artifacts'
 TRANSFORMED_COMPRESSED_PATH = 'artifacts/training_transformer.pkl'
 PREDICTION_COMPRESSED_PATH = 'artifacts/stacking_model.pkl'
@@ -56,8 +62,16 @@ def force_memory_release(logger: Logger):
         logger.info(f'Memory release failed: {e}')
 
 
-class MLFlowRepository:
-    def __init__(self, host: str, username: str, password: str, logger: Logger):
+class MLFlowRepository(SientiaMonitoring):
+    def __init__(
+        self,
+        host: str,
+        username: str,
+        password: str,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
+    ):
         """Initialize MLflow client and base state.
 
         Args:
@@ -67,6 +81,7 @@ class MLFlowRepository:
             logger (Logger): Logger instance.
         """
         # set tracking uri
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
         mlflow.set_tracking_uri(host)
 
         environ['MLFLOW_TRACKING_USERNAME'] = username
@@ -204,12 +219,15 @@ class MLFlowRepository:
     Functions related to download and load models
     """
 
-    def dowload_artifacts(self, model_name: str, artifact_path: str = 'data_model') -> str:
+    async def dowload_artifacts(
+        self, model_name: str, metadata: dict[str, Any], artifact_path: str = 'data_model'
+    ) -> str:
         """
         Download artifacts from the latest production run of a model.
 
         Args:
             model_name (str): Registered model name.
+            metadata (dict[str, Any]): Metadata used for structured logging.
             artifact_path (str): Relative path to artifacts within the run.
 
         Returns:
@@ -227,14 +245,29 @@ class MLFlowRepository:
 
         self.logger.info(f'Downloading artifacts from {run_id} to {output_dir}')
 
-        return self.client.download_artifacts(run_id, artifact_path, output_dir)
+        core_labels = self.get_core_labels(metadata, operation_type='download_artifacts')
 
-    def load_predict_model(self, model_name: str, flavor: str = 'sklearn') -> Any:
+        start_time = time.time()
+        try:
+            artifacts = self.client.download_artifacts(run_id, artifact_path, output_dir)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
+
+        return artifacts
+
+    async def load_predict_model(
+        self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
+    ) -> Any:
         """
         Load a predictive model from the MLflow Model Registry.
 
         Args:
             model_name (str): The name of the model to download from the registry.
+            metadata (dict[str, Any]): Metadata used for structured logging.
             flavor (str): Model flavor ('pyfunc', 'sklearn', 'pytorch')
             artifact_path (str | None): Path to compressed artifacts if model is compressed
 
@@ -247,18 +280,29 @@ class MLFlowRepository:
         """
         model_uri = f'models:/{model_name}/production'
         self.logger.info(f'Loading prediction model {model_name} from {model_uri}')
-        if flavor == 'pyfunc':
-            model = mlflow.pyfunc.load_model(model_uri)
-        elif flavor == 'sklearn':
-            model = mlflow.sklearn.load_model(model_uri)
-        elif flavor == 'pytorch':
-            model = mlflow.pytorch.load_model(model_uri)
-        else:
-            raise ValueError(INVALID_FLAVOR_MESSAGE)
 
+        core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
+        start_time = time.time()
+        try:
+            if flavor == 'pyfunc':
+                model = mlflow.pyfunc.load_model(model_uri)
+            elif flavor == 'sklearn':
+                model = mlflow.sklearn.load_model(model_uri)
+            elif flavor == 'pytorch':
+                model = mlflow.pytorch.load_model(model_uri)
+            else:
+                raise ValueError(INVALID_FLAVOR_MESSAGE)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
         return model
 
-    def load_transform_model(self, model_name: str, flavor: str) -> Any:
+    async def load_transform_model(
+        self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn'
+    ) -> Any:
         """
         Load the latest Production version of a transformation model.
 
@@ -267,6 +311,7 @@ class MLFlowRepository:
 
         Args:
             model_name (str): The name of the model to download.
+            metadata (dict[str, Any]): Metadata used for structured logging.
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
             artifact_path (str | None): Path to compressed artifacts if model is compressed
 
@@ -282,24 +327,41 @@ class MLFlowRepository:
         model_uri = self.get_model_uri(latest_production_id, prediction=False)
 
         self.logger.info(f'Loading data model {model_name} from {model_uri}')
-        if flavor == 'sklearn':
-            model = mlflow.sklearn.load_model(model_uri)
-        elif flavor == 'pyfunc':
-            model = mlflow.pyfunc.load_model(model_uri)
-        elif flavor == 'pytorch':
-            model = mlflow.pytorch.load_model(model_uri)
-        else:
-            raise ValueError(INVALID_FLAVOR_MESSAGE)
+
+        core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
+        start_time = time.time()
+        try:
+            if flavor == 'sklearn':
+                model = mlflow.sklearn.load_model(model_uri)
+            elif flavor == 'pyfunc':
+                model = mlflow.pyfunc.load_model(model_uri)
+            elif flavor == 'pytorch':
+                model = mlflow.pytorch.load_model(model_uri)
+            else:
+                raise ValueError(INVALID_FLAVOR_MESSAGE)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels)
+
         return model
 
-    def download_model(
-        self, model_name: str, model_type: str, flavor: str, load_wrapper: bool = False
+    async def download_model(
+        self,
+        model_name: str,
+        metadata: dict[str, Any],
+        model_type: str,
+        flavor: str,
+        load_wrapper: bool = False,
     ) -> tuple[Any, str | None]:
         """
         Download model based on type ("predict" or "transform").
 
         Args:
             model_name (str): Name of the model to download
+            metadata (dict[str, Any]): Metadata used for structured logging.
             model_type (str): Type of model ('predict' or 'transform')
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
             load_wrapper (bool): Whether to load wrapper
@@ -324,7 +386,7 @@ class MLFlowRepository:
 
             target = 'prediction_model' if model_type == 'predict' else 'data_model'
 
-            artifact_path = self.dowload_artifacts(model_name, target)
+            artifact_path = await self.dowload_artifacts(model_name, metadata, target)
 
             self.logger.info(
                 f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
@@ -334,10 +396,10 @@ class MLFlowRepository:
             model = raw_model._model_impl.python_model
         else:
             if model_type == 'predict':
-                model = self.load_predict_model(model_name, flavor)
+                model = await self.load_predict_model(model_name, metadata, flavor)
 
             else:
-                model = self.load_transform_model(model_name, flavor)
+                model = await self.load_transform_model(model_name, metadata, flavor)
 
         return model, artifact_path
 
@@ -446,12 +508,20 @@ class MLFlowRepository:
         del self.model_cache[model_key]['target']
         del self.model_cache[model_key]
 
-    def get_model(self, model_name: str, retention: int, model_type: str, flavor: str) -> Any:
+    async def get_model(
+        self,
+        model_name: str,
+        metadata: dict[str, Any],
+        retention: int,
+        model_type: str,
+        flavor: str,
+    ) -> Any:
         """
         Retrieve a model with caching support based on retention policy.
 
         Args:
             model_name (str): Name of the model to retrieve
+            metadata (dict[str, Any]): Metadata used for structured logging.
             retention (int): Cache retention time in minutes (0 = no cache).
             model_type (str): Type of model ('predict' or 'transform')
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch')
@@ -461,8 +531,12 @@ class MLFlowRepository:
         """
         # Retention is 0, download a new model
         if retention <= 0:
-            model, _artifact_path = self.download_model(
-                model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False
+            model, _artifact_path = await self.download_model(
+                model_name=model_name,
+                metadata=metadata,
+                model_type=model_type,
+                flavor=flavor,
+                load_wrapper=False,
             )
             return model
 
@@ -485,8 +559,12 @@ class MLFlowRepository:
                 )
 
         # Donwload new model (without lock to avoid blocking other threads)
-        model, _artifact_path = self.download_model(
-            model_name=model_name, model_type=model_type, flavor=flavor, load_wrapper=False
+        model, _artifact_path = await self.download_model(
+            model_name=model_name,
+            metadata=metadata,
+            model_type=model_type,
+            flavor=flavor,
+            load_wrapper=False,
         )
 
         # Update cache with lock
@@ -497,27 +575,35 @@ class MLFlowRepository:
         return model
 
     @overload
-    def get_cached_operation(
+    async def get_cached_operation(
         self,
         model_name: str,
         data: pd.DataFrame,
         operation: Literal['transform'],
         retention: int,
         flavor: str,
+        metadata: dict[str, Any],
     ) -> pd.DataFrame: ...
 
     @overload
-    def get_cached_operation(
+    async def get_cached_operation(
         self,
         model_name: str,
         data: pd.DataFrame,
         operation: Literal['predict'],
         retention: int,
         flavor: str,
+        metadata: dict[str, Any],
     ) -> pd.DataFrame | ndarray: ...
 
-    def get_cached_operation(
-        self, model_name: str, data: pd.DataFrame, operation: str, retention: int, flavor: str
+    async def get_cached_operation(
+        self,
+        model_name: str,
+        data: pd.DataFrame,
+        operation: str,
+        retention: int,
+        flavor: str,
+        metadata: dict[str, Any],
     ) -> pd.DataFrame | ndarray:
         """
         Execute a cached operation using the requested model.
@@ -527,15 +613,19 @@ class MLFlowRepository:
             data (pd.DataFrame): Input data.
             retention (int): Cache retention in minutes.
             flavor (str): Model flavor ('sklearn', 'pyfunc', 'pytorch').
-
+            metadata (dict[str, Any]): Metadata used for structured logging.
         Returns:
             pd.DataFrame | ndarray: Operation result.
         """
         if operation not in ['transform', 'predict']:
             raise ValueError("Invalid operation. Use 'transform' or 'predict'.")
 
-        model = self.get_model(
-            model_name=model_name, retention=retention, model_type=operation, flavor=flavor
+        model = await self.get_model(
+            model_name=model_name,
+            metadata=metadata,
+            retention=retention,
+            model_type=operation,
+            flavor=flavor,
         )
 
         prediction = model.predict(data)
@@ -552,7 +642,7 @@ class MLFlowRepository:
     Functions related to model retraining
     """
 
-    def fit_models(
+    async def fit_models(
         self,
         model_name: str,
         data: pd.DataFrame,
@@ -601,8 +691,9 @@ class MLFlowRepository:
 
         load_transform_wrapper = transform_flavor == 'pyfunc'
 
-        data_model, data_artifact_path = self.download_model(
+        data_model, data_artifact_path = await self.download_model(
             model_name=model_name,
+            metadata=metadata,
             model_type='transform',
             flavor=transform_flavor,
             load_wrapper=load_transform_wrapper,
@@ -612,8 +703,9 @@ class MLFlowRepository:
 
         load_predict_wrapper = predict_flavor == 'pyfunc'
 
-        prediction_model, prediction_artifact_path = self.download_model(
+        prediction_model, prediction_artifact_path = await self.download_model(
             model_name=model_name,
+            metadata=metadata,
             model_type='predict',
             flavor=predict_flavor,
             load_wrapper=load_predict_wrapper,
@@ -688,7 +780,7 @@ class MLFlowRepository:
         }
         return retrain_data
 
-    def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict):
+    async def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict):
         """Log a model into the active MLflow run.
 
         Args:
@@ -700,22 +792,34 @@ class MLFlowRepository:
         model = model_data['model']
 
         self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata)
-        if flavor == 'sklearn':
-            mlflow.sklearn.log_model(model, model_type)
-        elif flavor == 'pyfunc':
-            code_path = [path.join(model_data['artifact_path'], 'code', 'utils')]
 
-            self.logger.custom_debug(f'Code path: {code_path}', metadata)
+        core_labels = self.get_core_labels(metadata, operation_type='log_model')
+        start_time = time.time()
 
-            model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False)
+        try:
+            if flavor == 'sklearn':
+                mlflow.sklearn.log_model(model, model_type)
+            elif flavor == 'pyfunc':
+                code_path = [path.join(model_data['artifact_path'], 'code', 'utils')]
 
-            self.logger.custom_debug('Model uploaded successfully', metadata)
-        elif flavor == 'pytorch':
-            mlflow.pytorch.log_model(model, model_type)
-        else:
-            raise ValueError(INVALID_FLAVOR_MESSAGE)
+                self.logger.custom_debug(f'Code path: {code_path}', metadata)
 
-    def create_new_experiment(
+                model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False)
+
+                self.logger.custom_debug('Model uploaded successfully', metadata)
+            elif flavor == 'pytorch':
+                mlflow.pytorch.log_model(model, model_type)
+            else:
+                raise ValueError(INVALID_FLAVOR_MESSAGE)
+
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels)
+
+    async def create_new_experiment(
         self,
         model_name: str,
         data: pd.DataFrame,
@@ -784,30 +888,40 @@ class MLFlowRepository:
             metadata,
         )
 
-        with mlflow.start_run(
-            experiment_id=experiment.experiment_id,
-            run_name=current_run_name,
-            description=experiment_description,
-        ) as _run:
-            run_id = _run.info.run_id
-            self.logger.custom_info('Logging data model', metadata)
-            # dynamic parameters, including model itself
-            self.log_model(data_model, transform_flavor, 'data_model', metadata)
+        core_labels = self.get_core_labels(metadata, operation_type='create_new_experiment')
+        start_time = time.time()
+        try:
+            with mlflow.start_run(
+                experiment_id=experiment.experiment_id,
+                run_name=current_run_name,
+                description=experiment_description,
+            ) as _run:
+                run_id = _run.info.run_id
+                self.logger.custom_info('Logging data model', metadata)
+                # dynamic parameters, including model itself
+                await self.log_model(data_model, transform_flavor, 'data_model', metadata)
 
-            # dynamic parameters, including model itself
-            self.logger.custom_info('Logging prediction model', metadata)
-            self.log_model(prediction_model, predict_flavor, 'prediction_model', metadata)
+                # dynamic parameters, including model itself
+                self.logger.custom_info('Logging prediction model', metadata)
+                await self.log_model(prediction_model, predict_flavor, 'prediction_model', metadata)
 
-            self.logger.custom_info(f'Model logged successfully for {model_name}', metadata)
+                self.logger.custom_info(f'Model logged successfully for {model_name}', metadata)
 
-            self.logger.custom_info(f'Logging remaining parameters for {model_name}', metadata)
+                self.logger.custom_info(f'Logging remaining parameters for {model_name}', metadata)
 
-            # update transfomation model
-            # fixed parameters
-            mlflow.log_params(retrain_params)
+                # update transfomation model
+                # fixed parameters
+                mlflow.log_params(retrain_params)
 
-            # log the data raw
-            mlflow.log_artifact(data_path)
+                # log the data raw
+                mlflow.log_artifact(data_path)
+
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels)
 
         self.logger.custom_info('Deleting model from filesystem', metadata)
         if path.exists(model_temp_path):
@@ -829,7 +943,7 @@ class MLFlowRepository:
             'experiment_name': experiment.name,
         }
 
-    def update_production_model_by_run_id(
+    async def update_production_model_by_run_id(
         self, run_id: str, model_name: str, metadata: dict
     ) -> dict:
         """
@@ -864,7 +978,17 @@ class MLFlowRepository:
         # Registrar o modelo
         # Aqui estamos assumindo que você já tem um modelo salvo, caso contrário você precisará treiná-lo e salvá-lo primeiro.
         # Se o modelo já está registrado, você pode usar o método register_model() ou pyfunc.load_model() para isso.
-        mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
+
+        core_labels = self.get_core_labels(metadata, operation_type='register_model')
+        start_time = time.time()
+        try:
+            mlflow.register_model(f'runs:/{run_id}/prediction_model', model_name)
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels)
 
         # Obter a versão mais recente registrada do modelo
         model_versions = self.client.get_registered_model(model_name).latest_versions
@@ -875,9 +999,23 @@ class MLFlowRepository:
         max_version = max(model_versions, key=lambda x: int(x.version)).version
 
         # Mover a versão mais recente do modelo para o estágio de 'Production'
-        self.client.transition_model_version_stage(
-            name=model_name, version=max_version, stage='Production', archive_existing_versions=True
+        core_labels = self.get_core_labels(
+            metadata, operation_type='transition_model_version_stage'
         )
+        start_time = time.time()
+        try:
+            self.client.transition_model_version_stage(
+                name=model_name,
+                version=max_version,
+                stage='Production',
+                archive_existing_versions=True,
+            )
+        except Exception as e:
+            await self.emit_metric(metric_object=metrics.MODEL_WRITE_ERROR_COUNT, tags=core_labels)
+            raise e
+
+        await self.observe_lag(start_time, metrics.MODEL_WRITE_LAG, core_labels)
+        await self.emit_metric(metric_object=metrics.MODEL_WRITE_COUNT, tags=core_labels)
 
         return {'model_name': model_name, 'version': max_version, 'mlflow_run_id': run_id}
 
@@ -885,7 +1023,9 @@ class MLFlowRepository:
     Functions that provide the interface to model operations
     """
 
-    def transform(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict):
+    async def transform(
+        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
+    ):
         """
         Transform data using a cached transformation model.
 
@@ -930,8 +1070,13 @@ class MLFlowRepository:
         flavor = model_config.get('transform_flavor', 'sklearn')
 
         try:
-            transformed_data: pd.DataFrame = self.get_cached_operation(
-                model_name, data, 'transform', model_retention, flavor
+            transformed_data: pd.DataFrame = await self.get_cached_operation(
+                model_name=model_name,
+                data=data,
+                operation='transform',
+                retention=model_retention,
+                flavor=flavor,
+                metadata=metadata,
             )
 
             self.logger.custom_debug(
@@ -952,7 +1097,9 @@ class MLFlowRepository:
                 'content': {'message': str(e), 'traceback': traceback.format_exc()},
             }
 
-    def predict(self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict):
+    async def predict(
+        self, model_name: str, data: pd.DataFrame, model_config: dict, metadata: dict
+    ):
         """
         Generate predictions using a cached prediction model.
 
@@ -1007,8 +1154,13 @@ class MLFlowRepository:
             # data.to_csv(
             #     f"tmp/treated_data_{model_name}.csv", index=True)
 
-            predict_data = self.get_cached_operation(
-                model_name, data, 'predict', model_retention, flavor
+            predict_data = await self.get_cached_operation(
+                model_name=model_name,
+                data=data,
+                operation='predict',
+                retention=model_retention,
+                flavor=flavor,
+                metadata=metadata,
             )
 
             end_time = datetime.now()
@@ -1038,7 +1190,7 @@ class MLFlowRepository:
                 'content': {'message': str(e), 'traceback': traceback.format_exc()},
             }
 
-    def retrain_model(
+    async def retrain_model(
         self, data: pd.DataFrame, model_name: str, model_config: dict, metadata: dict
     ) -> dict[str, Any]:
         """
@@ -1099,7 +1251,7 @@ class MLFlowRepository:
         try:
             latest_production_id = self.get_model_run_id(model_name, stage='Production')
             self.logger.custom_info('Creating model experiment environment', metadata)
-            retrain_data = self.fit_models(
+            retrain_data = await self.fit_models(
                 model_name=model_name,
                 data=data,
                 transform_flavor=transform_flavor,
@@ -1113,7 +1265,7 @@ class MLFlowRepository:
             )
 
             self.logger.custom_info('Saving model retrain', metadata)
-            experiment = self.create_new_experiment(
+            experiment = await self.create_new_experiment(
                 model_name=model_name,
                 data=data,
                 retrain_data=retrain_data,
@@ -1141,7 +1293,7 @@ class MLFlowRepository:
                 'traceback': traceback.format_exc(),
             }
 
-    def update_production_model(
+    async def update_production_model(
         self, experiment: dict[str, Any], model_name: str, metadata: dict
     ) -> dict:
         """
@@ -1190,7 +1342,7 @@ class MLFlowRepository:
         """
         run_id = experiment['run_id']
         experiment_id = experiment['experiment_id']
-        metadata_result = self.update_production_model_by_run_id(run_id, model_name, metadata)
+        metadata_result = await self.update_production_model_by_run_id(run_id, model_name, metadata)
 
         metadata_result['mlflow_experiment_id'] = experiment_id
 

laborious/utils/repository/opc_repository.py

@@ -8,11 +8,14 @@ from typing import Any
 
 from asyncua import Client
 from asyncua.crypto.security_policies import SecurityPolicyBasic256
-from asyncua.ua import DataValue, DateTime, Variant, VariantType
+from asyncua.ua import DataValue, Variant, VariantType
 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.temporal.activities.base import BaseActivity
+from sientia_do.observability.metrics_controller import MetricsController
+from sientia_do.observability.sientia_monitoring import SientiaMonitoring
+
+from laborious import metrics
 
 data_type_map = {
     'float': {
@@ -38,13 +41,14 @@ data_type_map = {
 }
 
 
-class OpcRepository(BaseActivity):
+class OpcRepository(SientiaMonitoring):
     def __init__(
         self,
         opc_id: str,
         url: str,
         logger: Logger,
         notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
         reconnection_interval: int = 60,
         server_uri: str | None = None,
         cert_path: str | None = None,
@@ -61,10 +65,11 @@ class OpcRepository(BaseActivity):
         self.error_count = 0
         self.reconnection_interval = reconnection_interval
         self.last_reconnection_time: None | datetime = None
+        self.disconnection_interval = 10.0
         self.notification_handler = notification_handler
         self.client: None | Client = None
 
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
 
         self.metadata = {
             'model_name': '-',
@@ -164,6 +169,17 @@ class OpcRepository(BaseActivity):
                     'level': NotificationLevel.ERROR,
                 }
             await self.client.connect()
+
+            await self.emit_metric(
+                metric_object=metrics.OPC_CONNECTION_STATUS,
+                method='set',
+                tags={
+                    'pod_id': self.pod_id,
+                    'opc_server_id': self.id,
+                },
+                value=1,
+            )
+
             return True, {}
         except Exception as e:
             self.disconnect()
@@ -202,7 +218,7 @@ class OpcRepository(BaseActivity):
                         'traceback': traceback.format_exc(),
                     }
                 )
-                await asyncio.sleep(0.1 * i)
+                await asyncio.sleep(self.disconnection_interval * i)
         return error_stack
 
     async def disconnect(self):
@@ -218,7 +234,7 @@ class OpcRepository(BaseActivity):
 
         errors = await self.disconnection_fallback()
         if errors:
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=self.metadata,
                 notification_id=f'OPC_DISCONNECTION_ERROR_{self.id}',
                 message='Failed to disconnect from OPC server in 5 attempts.',
@@ -228,6 +244,15 @@ class OpcRepository(BaseActivity):
             )
         else:
             self.logger.warning(f'Disconnected from OPC server {self.id} successfully')
+        await self.emit_metric(
+            metric_object=metrics.OPC_CONNECTION_STATUS,
+            method='set',
+            tags={
+                'pod_id': self.pod_id,
+                'opc_server_id': self.id,
+            },
+            value=0,
+        )
 
         self.client = None
 
@@ -382,12 +407,12 @@ class OpcRepository(BaseActivity):
 
         data = data_type_map[data_type]['converter'](value)
         logger.custom_info(f'Writing {data} - {type(data)} to {node}', metadata)
-        now = datetime.now()
+        # now = datetime.now()
         ua_data = DataValue(
             Variant(data, data_type_map[data_type]['opc_type']),
-            SourceTimestamp=DateTime(
-                now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond
-            ),
+            # SourceTimestamp=DateTime(
+            #     now.year, now.month, now.day, now.hour, now.minute, now.second, now.microsecond
+            # ),
         )
 
         try: