diff --git a/encrypt.py b/encrypt.py index 6aff252..ad990a6 100644 --- a/encrypt.py +++ b/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 diff --git a/laborious/activities/gates.py b/laborious/activities/gates.py index 8cf4afa..9ae035a 100644 --- a/laborious/activities/gates.py +++ b/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) diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 1951973..52b56d9 100644 --- a/laborious/activities/opc.py +++ b/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'], - message=error_data['message'], - block=error_data['block'], - level=error_data.get('level', NotificationLevel.ERROR), - attachment_content=error_data.get('attachment_content', None), + notification_id=info_data['notification_id'], + message=info_data['message'], + block=info_data['block'], + level=info_data.get('level', NotificationLevel.ERROR), + attachment_content=info_data.get('attachment_content', None), ) - return False - return True + return None + 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 - success = success and local_success + response_times[tag] = response_time 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 - success = success and local_success + response_times[tag] = response_time - 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] diff --git a/laborious/metrics.py b/laborious/metrics.py index 7826fac..10d1221 100644 --- a/laborious/metrics.py +++ b/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], ) diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 8ab8563..2137c53 100644 --- a/laborious/utils/repository/opc_repository.py +++ b/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, + } diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 72ff12d..c8bb403 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/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), ) diff --git a/tests/laborious/activities/test_gates.py b/tests/laborious/activities/test_gates.py index 8d8822e..9229001 100644 --- a/tests/laborious/activities/test_gates.py +++ b/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, + ) diff --git a/tests/laborious/activities/test_opc.py b/tests/laborious/activities/test_opc.py index a9bbe67..669f82b 100644 --- a/tests/laborious/activities/test_opc.py +++ b/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'} - opc.write_data.assert_has_calls( - [ - call( - server_id='server1', - tag='tag1', - data=0.75, - data_type='float', - tag_type='prediction', - metadata=metadata['metadata'], - ) - ] + assert output == { + 'data': {'data': 'data'}, + 'metrics': {'server1': {'tag1': 0.1, 'tag2': 0.2}}, + } + opc.manage_output_tags.assert_called_once_with( + 'server1', + input_data['opc_output_config']['server1'], + mock_dataframe.return_value, + metadata['metadata'], + True, ) - opc.write_data.assert_has_calls( - [ - call( - server_id='server1', - tag='tag2', - data=0.95, - data_type='float', - tag_type='confidence', - metadata=metadata['metadata'], - ) - ] + opc.process_confidence.assert_called_once_with( + mock_dataframe.return_value, + True, + metadata['metadata'], ) - assert opc.write_data.call_count == 2 @mark.asyncio diff --git a/tests/laborious/utils/repository/test_opc_repository.py b/tests/laborious/utils/repository/test_opc_repository.py index 2a47f13..20c7bd2 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/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(mock_metrics, opc_repository, mock_client): +async def test_write_data(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, {}) - - 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 - ) + assert result == (True, {'response_time': ANY}) @pytest.mark.asyncio diff --git a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py b/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py index 006ef4e..59e39b3 100644 --- a/tests/laborious/workflows/subworkflows/test_format_and_export_prediction.py +++ b/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