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/activities/activities.py

@@ -1,3 +1,4 @@
+from sientia_do.observability.metrics_controller import MetricsController
 from temporalio import workflow
 
 with workflow.unsafe.imports_passed_through():
@@ -62,6 +63,8 @@ class Activities(Storage, MLFlow, Gates, OPC):
         Raises:
             Exception: If any parent class initialization fails
         """
+        metrics_controller = MetricsController(logger=logger)
+
         # Initialize parent classes
         Storage.__init__(
             self,
@@ -75,6 +78,7 @@ class Activities(Storage, MLFlow, Gates, OPC):
             minio_config=minio_config,
             logger=logger,
             notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
         )
 
         MLFlow.__init__(
@@ -86,12 +90,22 @@ class Activities(Storage, MLFlow, Gates, OPC):
             minio_config=minio_config,
             logger=logger,
             notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
         )
 
-        Gates.__init__(self, logger=logger, notification_handler=notification_handler)
+        Gates.__init__(
+            self,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
+        )
 
         OPC.__init__(
-            self, opc_servers=opc_config, logger=logger, notification_handler=notification_handler
+            self,
+            opc_servers=opc_config,
+            logger=logger,
+            notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
         )
 
     async def shutdown(self):
@@ -107,4 +121,6 @@ class Activities(Storage, MLFlow, Gates, OPC):
         proper resource cleanup and prevent resource leaks.
         """
         Storage.close(self)
-        await OPC.shutdown(self)
+        MLFlow.close(self)
+        Gates.close(self)
+        await OPC.close(self)

laborious/activities/gates.py

@@ -10,7 +10,8 @@ with workflow.unsafe.imports_passed_through():
     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 sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ, now
 
     from laborious import metrics
@@ -62,7 +63,7 @@ mlflow_content_path_confidence: Mapping[str, int] = {
 }
 
 
-class Gates(BaseActivity):
+class Gates(SientiaMonitoring):
     """
     Data quality gates and filtering activities for the Laborious system.
 
@@ -82,7 +83,12 @@ class Gates(BaseActivity):
         mlflow_content_filter_functions (dict): Mapping of MLFlow content filter names to functions
     """
 
-    def __init__(self, logger: Logger, notification_handler: NotificationHandler):
+    def __init__(
+        self,
+        logger: Logger,
+        notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
+    ):
         """
         Initialize data quality gates with logging and notification capabilities.
 
@@ -93,7 +99,16 @@ class Gates(BaseActivity):
         Raises:
             Exception: If BaseActivity initialization fails
         """
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
+
+    def close(self) -> None:
+        """
+        Close the gates activity and clean up resources.
+        """
+        SientiaMonitoring.shutdown(self)
+
+    def __del__(self):
+        self.close()
 
     @activity.defn(name='input_gate')
     async def input_gate(self, input_data: dict[str, Any]) -> tuple[str | None, int, str]:
@@ -153,7 +168,7 @@ class Gates(BaseActivity):
                     filter_output.append(config['policy'])
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
                     notification_id=f'INTPUT_GATE_ERROR__{fil}',
                     message=f'Error in filter {fil}:{config}: \n {e}',
@@ -225,7 +240,7 @@ class Gates(BaseActivity):
                 if mlflow_response_filter_functions[fil](data, config):
                     filter_output.append(config['policy'])
                     comments.append(data['content']['message'])
-                    self.send_notification(
+                    await self.send_notification_async(
                         metadata=metadata,
                         notification_id=f'{gate_type.upper()}_GATE_RESPONSE_FILTER__{fil}',
                         message=data['content']['message'],
@@ -235,7 +250,7 @@ class Gates(BaseActivity):
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
                     notification_id=f'MLFLOW_GATE_RESPONSE_FILTER__{fil}',
                     message=f'Error in filter {fil}:{config}: \n {e}',
@@ -305,7 +320,7 @@ class Gates(BaseActivity):
             try:
                 if mlflow_content_filter_functions[fil](data, config):
                     filter_output.append(config['policy'])
-                    self.send_notification(
+                    await self.send_notification_async(
                         metadata=metadata,
                         notification_id=f'{gate_type.upper()}_GATE_CONTENT_FILTER__{fil}',
                         message=f'Data not passed the content filter {fil}:{config}',
@@ -315,7 +330,7 @@ class Gates(BaseActivity):
                     )
             except Exception as e:
                 trace = traceback.format_exc()
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
                     notification_id=f'MLFLOW_GATE_CONTENT_FILTER__{fil}',
                     message=f'Error in filter {fil}:{config}: \n {e}',
@@ -608,39 +623,61 @@ class Gates(BaseActivity):
 
         self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
 
-        metrics.PREDICTIONS_WRITTEN_COUNT.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name'],
-        ).inc()
+        await self.emit_metric(
+            metric_object=metrics.PREDICTIONS_WRITTEN_COUNT,
+            tags={
+                'pod_id': self.pod_id,
+                'model_name': metadata['model_name'],
+                'workflow_name': metadata['workflow_name'],
+            },
+        )
 
-        metrics.PREDICTION_CONFIDENCE_MONITOR.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name'],
-        ).set(prediction_confidence)
+        await self.emit_metric(
+            metric_object=metrics.PREDICTION_CONFIDENCE_MONITOR,
+            method='set',
+            tags={
+                'pod_id': self.pod_id,
+                'model_name': metadata['model_name'],
+                'workflow_name': metadata['workflow_name'],
+            },
+            value=prediction_confidence,
+        )
 
-        metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels(
-            pod_id=self.pod_id,
-            model_name=metadata['model_name'],
-            pipeline_name=metadata['workflow_name'],
-        ).observe(response_time)
+        await self.emit_metric(
+            metric_object=metrics.PREDICTION_RESPONSE_TIME_MONITOR,
+            method='observe',
+            tags={
+                'pod_id': self.pod_id,
+                'model_name': metadata['model_name'],
+                'workflow_name': metadata['workflow_name'],
+            },
+            value=response_time,
+        )
 
         for server_id, tags in opc_metrics.items():
             for tag, response_time in tags.items():
-                metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels(
-                    pod_id=self.pod_id,
-                    model_name=metadata['model_name'],
-                    pipeline_name=metadata['workflow_name'],
-                    opc_server_id=server_id,
-                    tag=tag,
-                ).observe(response_time)
-                metrics.PREDICTION_OPC_WRITING_COUNT.labels(
-                    pod_id=self.pod_id,
-                    model_name=metadata['model_name'],
-                    pipeline_name=metadata['workflow_name'],
-                    opc_server_id=server_id,
-                    tag=tag,
-                ).inc()
+                await self.emit_metric(
+                    metric_object=metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR,
+                    method='observe',
+                    tags={
+                        'pod_id': self.pod_id,
+                        'model_name': metadata['model_name'],
+                        'workflow_name': metadata['workflow_name'],
+                        'opc_server_id': server_id,
+                        'tag': tag,
+                    },
+                    value=response_time,
+                )
+
+                await self.emit_metric(
+                    metric_object=metrics.PREDICTION_OPC_WRITING_COUNT,
+                    tags={
+                        'pod_id': self.pod_id,
+                        'model_name': metadata['model_name'],
+                        'workflow_name': metadata['workflow_name'],
+                        'opc_server_id': server_id,
+                        'tag': tag,
+                    },
+                )
 
         self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

laborious/activities/mlflow.py

@@ -10,7 +10,8 @@ with workflow.unsafe.imports_passed_through():
     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 sientia_do.temporal.constants import (
         DATETIME_FORMAT,
         DATETIME_FORMAT_MS_WITH_TZ,
@@ -22,7 +23,7 @@ with workflow.unsafe.imports_passed_through():
     from laborious.utils.repository.model_repository import MLFlowRepository
 
 
-class MLFlow(BaseActivity):
+class MLFlow(SientiaMonitoring):
     """
     MLFlow integration activities for model inference operations.
 
@@ -50,6 +51,7 @@ class MLFlow(BaseActivity):
         mlflow_password: str,
         logger: Logger,
         notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
     ):
         """
         Initialize MLFlow activities with server configuration.
@@ -65,14 +67,19 @@ class MLFlow(BaseActivity):
         Raises:
             Exception: If MLFlowRepository initialization fails
         """
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
         self.mlflow_host = mlflow_host
         self.mlflow_port = mlflow_port
         self.mlflow_username = mlflow_username
         self.mlflow_password = mlflow_password
 
         self.model_monitoring_repository = MLFlowRepository(
-            f'{mlflow_host}:{mlflow_port}', mlflow_username, mlflow_password, logger
+            f'{mlflow_host}:{mlflow_port}',
+            mlflow_username,
+            mlflow_password,
+            logger,
+            notification_handler,
+            metrics_controller,
         )
 
         if not hasattr(self, 'minio_repository'):
@@ -87,8 +94,18 @@ class MLFlow(BaseActivity):
                 minio_secret_key=minio_config['secret_key'],
                 minio_region_name=minio_config['region_name'],
                 minio_default_bucket=minio_config['default_bucket'],
+                metrics_controller=metrics_controller,
             )
 
+    def close(self) -> None:
+        """
+        Close the MLFlow activity and clean up resources.
+        """
+        SientiaMonitoring.shutdown(self)
+
+    def __del__(self):
+        self.close()
+
     @activity.defn(name='request_transform')
     async def request_transform(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
@@ -143,7 +160,7 @@ class MLFlow(BaseActivity):
         self.debug(data.head(5).to_string(), metadata)
 
         # Request transformation from MLFlow model
-        response_data = self.model_monitoring_repository.transform(
+        response_data = await self.model_monitoring_repository.transform(
             model_name, data, model_config, metadata
         )
 
@@ -208,7 +225,7 @@ class MLFlow(BaseActivity):
         ).dt.strftime(DATETIME_FORMAT)
 
         # Request prediction from MLFlow model
-        response_data = self.model_monitoring_repository.predict(
+        response_data = await self.model_monitoring_repository.predict(
             model_name, data, model_config, metadata
         )
 
@@ -263,12 +280,12 @@ class MLFlow(BaseActivity):
         self.info(f'Loading retrain data from Key: {object_key}', metadata)
 
         try:
-            data = self.minio_repository.get_parquet_as_dataframe(
+            data = await self.minio_repository.get_parquet_as_dataframe(
                 object_key=object_key, metadata=metadata
             )
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='ERROR_LOADING_RETRAIN_DATA',
                 message=f'Error loading retrain data: {e}',
@@ -316,13 +333,13 @@ class MLFlow(BaseActivity):
 
         data.columns.name = None
 
-        retrain_output = self.model_monitoring_repository.retrain_model(
+        retrain_output = await self.model_monitoring_repository.retrain_model(
             data=data, model_name=model_name, model_config=model_config, metadata=metadata
         )
 
         if not retrain_output['success']:
             trace = retrain_output['traceback']
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='RETRAIN_MODEL_ERROR',
                 message=f'Error retraining model {model_name}: {retrain_output["message"]}',
@@ -379,7 +396,7 @@ class MLFlow(BaseActivity):
         )
 
         try:
-            response = self.model_monitoring_repository.update_production_model(
+            response = await self.model_monitoring_repository.update_production_model(
                 experiment=experiment, model_name=model_name, metadata=metadata
             )
 
@@ -388,7 +405,7 @@ class MLFlow(BaseActivity):
 
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
                 message=f'Error updating production model {model_name}: {e}',

laborious/activities/opc.py

@@ -8,14 +8,15 @@ with workflow.unsafe.imports_passed_through():
     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.utils.repository.opc_repository import OpcRepository
 
 OPC_WRITTING_ERROR_CONFIDENCE = 12
 
 
-class OPC(BaseActivity):
+class OPC(SientiaMonitoring):
     """
     OPC server integration activities for real-time data export.
 
@@ -39,12 +40,13 @@ class OPC(BaseActivity):
         opc_servers: dict[str, dict[str, Any]],
         logger: Logger,
         notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
     ):
         self.logger = logger
         self.notification_handler = notification_handler
         self.opc_servers = opc_servers
 
-        BaseActivity.__init__(self, logger, notification_handler, set_error_counter=True)
+        SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
 
         self.opc_repository: dict[str, OpcRepository] = {}
         self.opc_servers = opc_servers
@@ -85,10 +87,11 @@ class OPC(BaseActivity):
                 server_cert_path=server['server_cert_path'],
                 notification_handler=self.notification_handler,
                 reconnection_interval=server['reconnection_interval'],
+                metrics_controller=self.metrics_controller,
             )
             is_connected, error_data = await self.opc_repository[opc_id].connect()
             if not is_connected:
-                self.send_notification(
+                await self.send_notification_async(
                     metadata={
                         'model_id': '-',
                         'model_name': '-',
@@ -137,7 +140,7 @@ class OPC(BaseActivity):
                 tag, data, data_type, self.logger, metadata
             )
             if not is_success:
-                self.send_notification(
+                await self.send_notification_async(
                     metadata=metadata,
                     notification_id=info_data['notification_id'],
                     message=info_data['message'],
@@ -149,7 +152,7 @@ class OPC(BaseActivity):
             return info_data['response_time']
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id=f'WRITE_OPC_{tag_type.upper()}_ERROR',
                 message=f'Error writing data to OPC server: {e}',
@@ -159,7 +162,7 @@ class OPC(BaseActivity):
             )
             raise e
 
-    def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
+    async def validate_server(self, server_id: str, metadata: dict[str, Any]) -> bool:
         """
         Validate that an OPC server is available and configured for write operations.
 
@@ -182,7 +185,7 @@ class OPC(BaseActivity):
         """
         if self.opc_repository.get(server_id) is None:
             message = f'OPC server {server_id} not found to perform write operation.'
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='OPC_SERVER_NOT_FOUND',
                 message=message,
@@ -299,7 +302,7 @@ class OPC(BaseActivity):
         metrics: dict[str, dict[str, float | None]] = {}
 
         for server_id, config in opc_output_config.items():
-            if not self.validate_server(server_id, metadata):
+            if not await self.validate_server(server_id, metadata):
                 success = False
                 continue
 
@@ -358,7 +361,7 @@ class OPC(BaseActivity):
 
         return data.to_dict()
 
-    async def shutdown(self):
+    async def close(self):
         """
         Gracefully shutdown all OPC server connections and cleanup resources.
 

laborious/activities/storage.py

@@ -9,6 +9,7 @@ with workflow.unsafe.imports_passed_through():
     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.temporal.activities.postgres import Postgres
     from sientia_do.temporal.constants import DATETIME_FORMAT_FILENAME, now
 
@@ -33,6 +34,7 @@ class Storage(Postgres):
         minio_config: dict[str, Any],
         logger: Logger,
         notification_handler: NotificationHandler,
+        metrics_controller: MetricsController,
     ):
         super().__init__(
             host=host,
@@ -44,6 +46,7 @@ class Storage(Postgres):
             max_connections=max_connections,
             logger=logger,
             notification_handler=notification_handler,
+            metrics_controller=metrics_controller,
         )
 
         if not hasattr(self, 'minio_repository'):
@@ -58,6 +61,7 @@ class Storage(Postgres):
                 minio_secret_key=minio_config['secret_key'],
                 minio_region_name=minio_config['region_name'],
                 minio_default_bucket=minio_config['default_bucket'],
+                metrics_controller=metrics_controller,
             )
 
     @activity.defn(name='query_to_minio')
@@ -95,14 +99,14 @@ class Storage(Postgres):
             data = pd.DataFrame(data)
 
             # Write parquet to memory and upload via persistent client
-            self.minio_repository.store_dataframe_as_parquet(
+            await self.minio_repository.store_dataframe_as_parquet(
                 dataframe=data, uri=uri, object_name=object_name, metadata=metadata
             )
 
             return {'success': True, 'object_key': object_name, 'uri': uri}
         except Exception as e:
             trace = traceback.format_exc()
-            self.send_notification(
+            await self.send_notification_async(
                 metadata=metadata,
                 notification_id='ERROR_STORING_QUERY_TO_MINIO',
                 message=f'Error storing query to MinIO: {e}',

laborious/metrics.py

@@ -19,11 +19,12 @@ Key Metric Categories:
 Metric Labels:
 - pod_id: Kubernetes pod identifier for multi-instance deployments
 - model_name: Name of the ML model being used
-- pipeline_name: Name of the prediction pipeline
+- workflow_name: Name of the prediction pipeline
 - opc_server_id: Identifier for OPC server operations
 """
 
 from prometheus_client import Counter, Gauge, Histogram
+from sientia_do.observability.metrics import CORE_LABELS as SIENTIA_CORE_LABELS
 
 # Application health metric
 APP_UP = Gauge(
@@ -33,7 +34,7 @@ APP_UP = Gauge(
 )
 
 # Core labels used across multiple metrics
-CORE_LABELS = ['pod_id', 'model_name', 'pipeline_name']
+CORE_LABELS = ['pod_id', 'model_name', 'workflow_name']
 
 # Prediction operation metrics
 PREDICTIONS_WRITTEN_COUNT = Counter(
@@ -49,7 +50,7 @@ PREDICTION_CONFIDENCE_MONITOR = Gauge(
     CORE_LABELS,
 )
 
-# Performance monitoring metrics
+# Prediction total response time
 PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
     'laborious_prediction_response_time_monitor',
     'Current response time of each prediction',
@@ -57,7 +58,49 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
     buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )
 
-# OPC export metrics
+# ================== MinIO metrics ==================
+
+MINIO_READ_LAG = Histogram(
+    'laborious_minio_read_lag',
+    'Lag between the last write to MinIO and the last read from MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MINIO_WRITE_LAG = Histogram(
+    'laborious_minio_write_lag',
+    'Lag between the last write to MinIO and the last read from MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MINIO_READ_COUNT = Counter(
+    'laborious_minio_read_count',
+    'Number of reads from MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+)
+
+MINIO_WRITE_COUNT = Counter(
+    'laborious_minio_write_count',
+    'Number of writes to MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+)
+
+MINIO_READ_ERROR_COUNT = Counter(
+    'laborious_minio_read_error_count',
+    'Number of errors reading from MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+)
+
+MINIO_WRITE_ERROR_COUNT = Counter(
+    'laborious_minio_write_error_count',
+    'Number of errors writing to MinIO',
+    [*SIENTIA_CORE_LABELS, 'bucket_name', 'object_name'],
+)
+
+# ================== OPC metrics ==================
+
+
 PREDICTION_OPC_WRITING_COUNT = Counter(
     'laborious_prediction_opc_writing_count',
     'Number of predictions written to the OPC server',
@@ -70,3 +113,49 @@ PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
     [*CORE_LABELS, 'opc_server_id', 'tag'],
     buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
 )
+
+OPC_CONNECTION_STATUS = Gauge(
+    'laborious_opc_connection_status',
+    'Connection status with the OPC server (1=connected, 0=disconnected)',
+    ['pod_id', 'opc_server_id'],
+)
+
+# ================== Model metrics ==================
+
+MODEL_READ_LAG = Histogram(
+    'laborious_model_read_lag',
+    'Lag between the start and read of read operations',
+    SIENTIA_CORE_LABELS,
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MODEL_WRITE_LAG = Histogram(
+    'laborious_model_write_lag',
+    'Lag between the start and end of write operations',
+    SIENTIA_CORE_LABELS,
+    buckets=[0.01, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0, 10.0],
+)
+
+MODEL_READ_COUNT = Counter(
+    'laborious_model_read_count',
+    'Number of reads from the model',
+    SIENTIA_CORE_LABELS,
+)
+
+MODEL_WRITE_COUNT = Counter(
+    'laborious_model_write_count',
+    'Number of writes to the model',
+    SIENTIA_CORE_LABELS,
+)
+
+MODEL_READ_ERROR_COUNT = Counter(
+    'laborious_model_read_error_count',
+    'Number of errors reading from the model',
+    SIENTIA_CORE_LABELS,
+)
+
+MODEL_WRITE_ERROR_COUNT = Counter(
+    'laborious_model_write_error_count',
+    'Number of errors writing to the model',
+    SIENTIA_CORE_LABELS,
+)

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: