SIENTIAPDE-1314

Enhance OPC Metrics Handling and Refactor Write Operations

- Updated the OPC class to return response times for write operations, improving metrics tracking.
- Refactored the Gates activity to incorporate OPC metrics into the metrics writing process.
- Adjusted the manage_output_tags method in OpcRepository to return response times for each tag written.
- Modified tests to validate the new metrics structure and ensure correct behavior of the updated methods.

encrypt.py

@@ -1,7 +1,8 @@
 import os
 import argparse
 from pathspec import PathSpec
-import yaml
+import yaml # type: ignore
+from typing import Any
 
 '''
 Usage:
@@ -33,7 +34,7 @@ def encode_file_tree_to_yaml(directory, ignore_file, include_library):
     """Encode the file tree into a single YAML file."""
     ignore_patterns = load_ignore_patterns(
         ignore_file, include_library) if ignore_file else None
-    file_tree = {}
+    file_tree: dict[str, Any] = {}
 
     for root, dirs, files in os.walk(directory):
         # Skip ignored directories

laborious/activities/gates.py

@@ -604,6 +604,7 @@ class Gates(BaseActivity):
         prediction = DataFrame(input_data['prediction'])
         prediction_confidence = prediction['prediction_confidence'].values[0]
         response_time = prediction['response_time'].values[0]
+        opc_metrics = input_data['opc_metrics']
 
         self.info(f'Writing metrics for model {metadata["model_name"]}', metadata)
 
@@ -625,4 +626,21 @@ class Gates(BaseActivity):
             pipeline_name=metadata['workflow_name'],
         ).observe(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()
+
         self.info(f'Metrics written for model {metadata["model_name"]}', metadata)

laborious/activities/opc.py

@@ -112,7 +112,7 @@ class OPC(BaseActivity):
         data_type: str,
         tag_type: str,
         metadata: dict[str, Any],
-    ) -> bool:
+    ) -> float | None:
         """
         Write data to a specific OPC server tag with comprehensive error handling.
 
@@ -133,20 +133,20 @@ class OPC(BaseActivity):
         """
 
         try:
-            is_success, error_data = await self.opc_repository[server_id].write_data(
+            is_success, info_data = await self.opc_repository[server_id].write_data(
                 tag, data, data_type, self.logger, metadata
             )
             if not is_success:
                 self.send_notification(
                     metadata=metadata,
-                    notification_id=error_data['notification_id'],
+                    notification_id=info_data['notification_id'],
-                    message=error_data['message'],
+                    message=info_data['message'],
-                    block=error_data['block'],
+                    block=info_data['block'],
-                    level=error_data.get('level', NotificationLevel.ERROR),
+                    level=info_data.get('level', NotificationLevel.ERROR),
-                    attachment_content=error_data.get('attachment_content', None),
+                    attachment_content=info_data.get('attachment_content', None),
                 )
-                return False
+                return None
-            return True
+            return info_data['response_time']
         except Exception as e:
             trace = traceback.format_exc()
             self.send_notification(
@@ -200,7 +200,7 @@ class OPC(BaseActivity):
         data: DataFrame,
         metadata: dict[str, Any],
         success: bool,
-    ) -> tuple[bool, int]:
+    ) -> tuple[bool, dict[str, float | None]]:
         """
         Manage the writing of prediction and confidence data to OPC server tags.
 
@@ -228,10 +228,11 @@ class OPC(BaseActivity):
                 - total_tags_written: Count of successfully written tags
         """
 
-        count = 0
+        response_times: dict[str, float | None] = {}
+
         if 'prediction_tags' in config:
             for tag, tag_config in config['prediction_tags'].items():
-                local_success = await self.write_data(
+                response_time = await self.write_data(
                     server_id=server_id,
                     tag=tag,
                     data=data.head(1)['prediction'].values[0],
@@ -239,17 +240,16 @@ class OPC(BaseActivity):
                     tag_type='prediction',
                     metadata=metadata,
                 )
-                if local_success:
+                if response_time is not None:
                     self.info(
                         f'Prediction data written to OPC server {server_id} for tag {tag}.',
                         metadata,
                     )
-                    count += 1
+                response_times[tag] = response_time
-                success = success and local_success
 
         if 'confidence_tags' in config:
             for tag, tag_config in config['confidence_tags'].items():
-                local_success = await self.write_data(
+                response_time = await self.write_data(
                     server_id=server_id,
                     tag=tag,
                     data=data.head(1)['prediction_confidence'].values[0],
@@ -257,15 +257,16 @@ class OPC(BaseActivity):
                     tag_type='confidence',
                     metadata=metadata,
                 )
-                if local_success:
+                if response_time is not None:
                     self.info(
                         f'Confidence data written to OPC server {server_id} for tag {tag}.',
                         metadata,
                     )
-                    count += 1
+                response_times[tag] = response_time
-                success = success and local_success
 
-        return success, count
+        success = None not in response_times.values()
+
+        return success, response_times
 
     @activity.defn(name='write_opc_data')
     async def write_opc_data(self, input_data: dict[str, Any]) -> dict[Any, Any]:
@@ -294,14 +295,18 @@ class OPC(BaseActivity):
 
         success = True
 
+        metrics: dict[str, dict[str, float | None]] = {}
+
         for server_id, config in opc_output_config.items():
             if not self.validate_server(server_id, metadata):
                 success = False
                 continue
 
-            local_success, local_count = await self.manage_output_tags(
+            local_success, local_response_times = await self.manage_output_tags(
                 server_id, config, data, metadata, success
             )
+            metrics[server_id] = local_response_times
+            local_count = len(local_response_times)
             success = success and local_success
 
             self.info(
@@ -309,7 +314,10 @@ class OPC(BaseActivity):
                 metadata,
             )
 
-        return self.process_confidence(data, success, metadata)
+        return {
+            'data': self.process_confidence(data, success, metadata),
+            'metrics': metrics,
+        }
 
     def process_confidence(
         self, data: DataFrame, success: bool, metadata: dict[str, Any]

laborious/metrics.py

@@ -61,12 +61,12 @@ PREDICTION_RESPONSE_TIME_MONITOR = Histogram(
 PREDICTION_OPC_WRITING_COUNT = Counter(
     'laborious_prediction_opc_writing_count',
     'Number of predictions written to the OPC server',
-    [*CORE_LABELS, 'opc_server_id'],
+    [*CORE_LABELS, 'opc_server_id', 'tag'],
 )
 
 PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR = Histogram(
     'laborious_prediction_opc_writing_response_time_monitor',
     'Current response time of each prediction written to the OPC server',
-    [*CORE_LABELS, 'opc_server_id'],
+    [*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],
 )

laborious/utils/repository/opc_repository.py

@@ -14,8 +14,6 @@ from sientia_do.notifications.models import NotificationLevel
 from sientia_do.observability.logger import Logger
 from sientia_do.temporal.activities.base import BaseActivity
 
-from laborious import metrics
-
 data_type_map = {
     'float': {
         'converter': float,
@@ -395,21 +393,8 @@ class OpcRepository(BaseActivity):
         try:
             await node_obj.write_value(ua_data)
 
-            metrics.PREDICTION_OPC_WRITING_COUNT.labels(
-                pod_id=self.pod_id,
-                model_name=metadata['model_name'],
-                pipeline_name=metadata['workflow_name'],
-                opc_server_id=self.id,
-            ).inc()
-
             end_time = time.time()
             response_time = end_time - start_time
-            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=self.id,
-            ).observe(response_time)
 
         except Exception as e:
             trace = traceback.format_exc()
@@ -424,4 +409,6 @@ class OpcRepository(BaseActivity):
             }
         self.error_count = 0
 
-        return True, {}
+        return True, {
+            'response_time': response_time,
+        }

laborious/workflows/sub_workflows/format_and_export_prediction.py

@@ -103,9 +103,13 @@ class FormatAndExportPrediction:
             )
 
         # write to opc
-        prediction = await workflow.execute_activity_method(
+        prediction, opc_metrics = await workflow.execute_activity_method(
             Activities.write_opc_data,
-            {**metadata, 'opc_output_config': input_data['opc_output_config'], 'data': prediction},
+            {
+                'opc_output_config': input_data['opc_output_config'],
+                'data': prediction,
+                **metadata,
+            },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=60),
         )
@@ -126,7 +130,11 @@ class FormatAndExportPrediction:
 
         await workflow.execute_activity_method(
             Activities.write_metrics,
-            {**metadata, 'prediction': prediction},
+            {
+                **metadata,
+                'prediction': prediction,
+                'opc_metrics': opc_metrics,
+            },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=60),
         )

tests/laborious/activities/test_gates.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, MagicMock, patch
+from unittest.mock import ANY, MagicMock, call, patch
 
 from pytest import fixture, mark
 from sientia_do.notifications.models import NotificationLevel
@@ -636,6 +636,7 @@ async def test_write_metrics(mock_metrics, gates_activity):
             'prediction_confidence': [0.9, 0.8, 0.7],
             'response_time': [0.1, 0.2, 0.3],
         },
+        'opc_metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
     }
     await gates_activity.write_metrics(input_data)
     mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
@@ -660,3 +661,35 @@ async def test_write_metrics(mock_metrics, gates_activity):
     mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
         0.1
     )
+
+    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_has_calls(
+        [
+            call(
+                pod_id=gates_activity.pod_id,
+                model_name=metadata['metadata']['model_name'],
+                pipeline_name=metadata['metadata']['workflow_name'],
+                opc_server_id='server1',
+                tag='tag1',
+            )
+        ]
+    )
+    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_has_calls(
+        [
+            call(
+                pod_id=gates_activity.pod_id,
+                model_name=metadata['metadata']['model_name'],
+                pipeline_name=metadata['metadata']['workflow_name'],
+                opc_server_id='server1',
+                tag='tag1',
+            )
+        ]
+    )
+
+    assert mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.call_count == 2
+    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_has_calls(
+        [
+            call(0.1),
+            call(0.2),
+        ],
+        any_order=True,
+    )

tests/laborious/activities/test_opc.py

@@ -180,6 +180,8 @@ WRITE_DATA_CASES = [
 @mark.parametrize('tag,data_type,data', WRITE_DATA_CASES)
 @mark.asyncio
 async def test_write_data_success(opc, tag, data_type, data):
+    opc.opc_repository['server1'].write_data.return_value = (True, {'response_time': 0.1})
+
     result = await opc.write_data(
         server_id='server1',
         tag=tag,
@@ -188,7 +190,7 @@ async def test_write_data_success(opc, tag, data_type, data):
         tag_type='prediction',
         metadata=metadata,
     )
-    assert result is True
+    assert result == 0.1
     opc.opc_repository['server1'].write_data.assert_called_once_with(
         tag, data, data_type, opc.logger, metadata
     )
@@ -215,7 +217,7 @@ async def test_write_data_failed(opc):
         tag_type='prediction',
         metadata=metadata,
     )
-    assert result is False
+    assert result is None
 
     opc.send_notification.assert_called_once_with(
         metadata=metadata,
@@ -256,7 +258,8 @@ async def test_write_data_exception(opc):
 
 
 @mark.asyncio
-async def test_write_opc_data_success(opc):
+@patch('laborious.activities.opc.DataFrame')
+async def test_write_opc_data_success(mock_dataframe, opc):
     # Arrange
     input_data = {
         **metadata,
@@ -270,37 +273,28 @@ async def test_write_opc_data_success(opc):
     }
 
     # Act
-    opc.write_data = AsyncMock(return_value=True)
+    opc.manage_output_tags = AsyncMock(return_value=(True, {'tag1': 0.1, 'tag2': 0.2}))
+
     opc.process_confidence = MagicMock(return_value={'data': 'data'})
     output = await opc.write_opc_data(input_data)
 
     # Assert
-    assert output == {'data': 'data'}
+    assert output == {
-    opc.write_data.assert_has_calls(
+        'data': {'data': 'data'},
-        [
+        'metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}},
-            call(
+    }
-                server_id='server1',
+    opc.manage_output_tags.assert_called_once_with(
-                tag='tag1',
+        'server1',
-                data=0.75,
+        input_data['opc_output_config']['server1'],
-                data_type='float',
+        mock_dataframe.return_value,
-                tag_type='prediction',
+        metadata['metadata'],
-                metadata=metadata['metadata'],
+        True,
-            )
-        ]
     )
-    opc.write_data.assert_has_calls(
+    opc.process_confidence.assert_called_once_with(
-        [
+        mock_dataframe.return_value,
-            call(
+        True,
-                server_id='server1',
+        metadata['metadata'],
-                tag='tag2',
-                data=0.95,
-                data_type='float',
-                tag_type='confidence',
-                metadata=metadata['metadata'],
-            )
-        ]
     )
-    assert opc.write_data.call_count == 2
 
 
 @mark.asyncio

tests/laborious/utils/repository/test_opc_repository.py

@@ -335,7 +335,7 @@ async def test_write_data_validate_connection_do_nothing(opc_repository):
 
     opc_repository.validate_connection.assert_called_once()
     opc_repository.client.get_node.assert_called_once_with('ns=2;s=TestNode')
-    assert result == (True, {})
+    assert result == (True, {'response_time': ANY})
 
 
 @pytest.mark.asyncio
@@ -403,8 +403,7 @@ async def test_write_data_invalid_data_type(opc_repository, mock_client):
 
 
 @pytest.mark.asyncio
-@patch('laborious.utils.repository.opc_repository.metrics')
+async def test_write_data(opc_repository, mock_client):
-async def test_write_data(mock_metrics, opc_repository, mock_client):
     opc_repository.validate_connection = AsyncMock(return_value=(True, {}))
     opc_repository.client = mock_client
     mock_node = AsyncMock()
@@ -416,25 +415,7 @@ async def test_write_data(mock_metrics, opc_repository, mock_client):
 
     mock_client.get_node.assert_called_once_with('ns=2;s=TestNode')
     mock_node.write_value.assert_called_once()
-    assert result == (True, {})
+    assert result == (True, {'response_time': ANY})
-
-    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.assert_called_once_with(
-        pod_id=opc_repository.pod_id,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id,
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_COUNT.labels.return_value.inc.assert_called_once_with()
-
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
-        pod_id=opc_repository.pod_id,
-        model_name=metadata['metadata']['model_name'],
-        pipeline_name=metadata['metadata']['workflow_name'],
-        opc_server_id=opc_repository.id,
-    )
-    mock_metrics.PREDICTION_OPC_WRITING_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
-        ANY
-    )
 
 
 @pytest.mark.asyncio

tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py

@@ -1,4 +1,4 @@
-from unittest.mock import ANY, AsyncMock, call, patch
+from unittest.mock import ANY, AsyncMock, MagicMock, call, patch
 
 from pytest import fixture, mark
 from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ
@@ -42,6 +42,15 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
         'prediction_store_policy': 'erl:1',
     }
 
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
     await format_and_export_prediction.run(input_data)
 
     workflow_mock.execute_local_activity_method.assert_has_calls(
@@ -84,7 +93,7 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
                 {
                     'schema': input_data['schema'],
                     'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': workflow_mock.execute_local_activity_method.return_value,
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
@@ -97,6 +106,21 @@ async def test_run_none_path_flag(workflow_mock, format_and_export_prediction):
         ]
     )
 
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1
 
@@ -121,6 +145,15 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
         'comment': 'test_comment',
     }
 
+    prediction_data = MagicMock()
+    opc_metrics = MagicMock()
+
+    workflow_mock.execute_activity_method.side_effect = [
+        (prediction_data, opc_metrics),
+        MagicMock(),
+        MagicMock(),
+    ]
+
     await format_and_export_prediction.run(input_data)
 
     workflow_mock.execute_local_activity_method.assert_has_calls(
@@ -162,7 +195,7 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
                 {
                     'schema': input_data['schema'],
                     'table_name': input_data['table_name'],
-                    'data': workflow_mock.execute_activity_method.return_value,
+                    'data': prediction_data,
                     **metadata,
                     'timestamp_conversion': {
                         'column': 'timestamp',
@@ -175,5 +208,20 @@ async def test_run_default_path_flag(workflow_mock, format_and_export_prediction
         ]
     )
 
+    workflow_mock.execute_activity_method.assert_has_calls(
+        [
+            call(
+                Activities.write_metrics,
+                {
+                    **metadata,
+                    'prediction': prediction_data,
+                    'opc_metrics': opc_metrics,
+                },
+                retry_policy=ANY,
+                start_to_close_timeout=ANY,
+            )
+        ]
+    )
+
     assert workflow_mock.execute_activity_method.call_count == 3
     assert workflow_mock.execute_local_activity_method.call_count == 1