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..d614433 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( @@ -199,8 +199,7 @@ class OPC(BaseActivity): config: dict[str, Any], 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 +227,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 +239,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,18 +256,21 @@ 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]: + async def write_opc_data( + self, input_data: dict[str, Any] + ) -> tuple[dict[Any, Any], dict[str, dict[str, float | None]]]: """ Write prediction and confidence data to OPC servers. The two writing operations are optional and independent of each other. @@ -294,14 +296,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( - server_id, config, data, metadata, success + local_success, local_response_times = await self.manage_output_tags( + server_id, config, data, metadata ) + metrics[server_id] = local_response_times + local_count = len(local_response_times) success = success and local_success self.info( @@ -309,7 +315,7 @@ class OPC(BaseActivity): metadata, ) - return self.process_confidence(data, success, metadata) + return self.process_confidence(data, success, metadata), 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/model_repository.py b/laborious/utils/repository/model_repository.py index 7400502..c585d7d 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,6 +16,7 @@ Capabilities: import ctypes import gc +import threading import traceback from datetime import datetime, timedelta from os import environ, makedirs, path @@ -74,6 +75,7 @@ class MLFlowRepository: self.client = mlflow.tracking.MlflowClient() self.model_cache: dict[str, Any] = {} + self._cache_lock = threading.RLock() self.logger = logger """ @@ -441,7 +443,7 @@ class MLFlowRepository: """ self.logger.debug(f'Model {model_name} is outdated, downloading a new one') - del self.model_cache[model_key]['target']['model'] + 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: @@ -466,28 +468,31 @@ class MLFlowRepository: model_key = f'{model_name}_{model_type}' - if model_key in self.model_cache: - cache = self.model_cache[model_key] + # Acquire lock to check cache + with self._cache_lock: + if model_key in self.model_cache: + cache = self.model_cache[model_key] - # Check if config has changed or is outdated - if self.check_cache_retention(cache, retention): - return self.handle_valid_model(model_name=model_name, cache=cache) + # Check if config has changed or is outdated + if self.check_cache_retention(cache, retention): + return self.handle_valid_model(model_name=model_name, cache=cache) + else: + # Model is outdated, delete old model files + self.handle_outdated_model(model_name=model_name, model_key=model_key) else: - # Model is outdated, delete old model files - self.handle_outdated_model(model_name=model_name, model_key=model_key) - else: - self.logger.debug( - f'Model {model_name} is not in {model_type} cache, downloading a new one' - ) + self.logger.debug( + f'Model {model_name} is not in {model_type} cache, downloading a new one' + ) - # Donwload new model + # 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 ) - cache = {'target': model, 'timestamp': datetime.now()} - - self.model_cache[model_key] = cache + # Update cache with lock + with self._cache_lock: + cache = {'target': model, 'timestamp': datetime.now()} + self.model_cache[model_key] = cache return model diff --git a/laborious/utils/repository/opc_repository.py b/laborious/utils/repository/opc_repository.py index 8ab8563..95bfd77 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, @@ -265,22 +263,22 @@ class OpcRepository(BaseActivity): if self.client is None: return await self.connect() - if self.error_count > 5: - self.logger.custom_warning( - f'OPC server {self.id} will be disconnected due to multiple errors', self.metadata - ) - try: - await self.disconnect() - except Exception as e: - trace = traceback.format_exc() - self.logger.custom_error( - f'Failed to disconnect from OPC server: {e}', self.metadata - ) - self.logger.custom_error(trace, self.metadata) - self.logger.custom_info( - f'Attempting to reconnect to OPC server {self.id}...', self.metadata - ) - return await self.connect() + # if self.error_count > 5: # NOSONAR + # self.logger.custom_warning( + # f'OPC server {self.id} will be disconnected due to multiple errors', self.metadata + # ) + # try: + # await self.disconnect() + # except Exception as e: + # trace = traceback.format_exc() + # self.logger.custom_error( + # f'Failed to disconnect from OPC server: {e}', self.metadata + # ) + # self.logger.custom_error(trace, self.metadata) + # self.logger.custom_info( + # f'Attempting to reconnect to OPC server {self.id}...', self.metadata + # ) + # return await self.connect() # Check if client is connected using asyncua's connection state try: @@ -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/worker/worker.py b/laborious/worker/worker.py index 781ad55..5faf004 100644 --- a/laborious/worker/worker.py +++ b/laborious/worker/worker.py @@ -152,7 +152,7 @@ async def main(): max_concurrent_workflow_tasks=50, max_concurrent_activities=50, max_concurrent_local_activities=50, - max_cached_workflows=200, + max_cached_workflows=2, workflow_task_poller_behavior=PollerBehaviorAutoscaling(), activity_task_poller_behavior=PollerBehaviorAutoscaling(), ), @@ -164,7 +164,6 @@ async def main(): # MLFlow activities.request_predict, activities.request_transform, - activities.query_to_minio, # Gates activities.input_gate, activities.mlflow_response_gate, diff --git a/laborious/workflows/predictions_batch.py b/laborious/workflows/predictions_batch.py index f79206f..afd7c86 100644 --- a/laborious/workflows/predictions_batch.py +++ b/laborious/workflows/predictions_batch.py @@ -114,4 +114,4 @@ class PredictionsBatch: } # Execute prediction process workflow - await workflow.execute_child_workflow('prediction_process', prediction_input) + await workflow.execute_child_workflow('subworkflow.prediction_process', prediction_input) diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py index 72ff12d..8d8c9da 100644 --- a/laborious/workflows/sub_workflows/format_and_export_prediction.py +++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py @@ -10,7 +10,7 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.activities import Activities -@workflow.defn(name='format_and_export_prediction') +@workflow.defn(name='subworkflow.format_and_export_prediction') class FormatAndExportPrediction: """ Data formatting and export workflow for prediction results. @@ -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), ) @@ -121,12 +125,16 @@ class FormatAndExportPrediction: 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ}, }, retry_policy=retry_policy, - start_to_close_timeout=timedelta(seconds=60), + start_to_close_timeout=timedelta(seconds=180), ) 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/laborious/workflows/sub_workflows/prediction_process.py b/laborious/workflows/sub_workflows/prediction_process.py index 7bd36e2..ab78fb1 100644 --- a/laborious/workflows/sub_workflows/prediction_process.py +++ b/laborious/workflows/sub_workflows/prediction_process.py @@ -9,7 +9,7 @@ with workflow.unsafe.imports_passed_through(): from laborious.activities.activities import Activities -@workflow.defn(name='prediction_process') +@workflow.defn(name='subworkflow.prediction_process') class PredictionProcess: """ Core prediction processing workflow for the Laborious system. @@ -194,7 +194,7 @@ class PredictionProcess: # Delegate to export workflow for data persistence await workflow.execute_child_workflow( - 'format_and_export_prediction', + 'subworkflow.format_and_export_prediction', { 'metadata': metadata, 'path_flag': path_flag, @@ -275,7 +275,7 @@ class PredictionProcess: elif path_flag == 'CONTINUE': # call write workflow await workflow.execute_child_workflow( - 'format_and_export_prediction', + 'subworkflow.format_and_export_prediction', { 'metadata': metadata, 'path_flag': path_flag, 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..4d1a86a 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,106 @@ async def test_write_data_exception(opc): @mark.asyncio -async def test_write_opc_data_success(opc): +async def test_manage_output_tags_success(opc): + opc.write_data = AsyncMock(return_value=0.1) + + data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]}) + config = { + 'prediction_tags': {'tag1': {'data_type': 'float'}}, + 'confidence_tags': {'tag2': {'data_type': 'float'}}, + } + + output_data, opc_metrics = await opc.manage_output_tags( + server_id='server1', + config=config, + data=data, + metadata=metadata['metadata'], + ) + + assert output_data is True + assert opc_metrics == {'tag1': 0.1, 'tag2': 0.1} + opc.write_data.assert_has_calls( + [ + call( + server_id='server1', + tag='tag1', + data=0.75, + data_type='float', + tag_type='prediction', + metadata=metadata['metadata'], + ), + call( + server_id='server1', + tag='tag2', + data=0.95, + data_type='float', + tag_type='confidence', + metadata=metadata['metadata'], + ), + ] + ) + + +@mark.asyncio +@mark.parametrize('side_effect', [[0.1, None], [None, 0.2]]) +async def test_manage_output_tags_failed(opc, side_effect): + opc.write_data = AsyncMock(side_effect=side_effect) + data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]}) + config = { + 'prediction_tags': {'tag1': {'data_type': 'float'}}, + 'confidence_tags': {'tag2': {'data_type': 'float'}}, + } + output_data, opc_metrics = await opc.manage_output_tags( + server_id='server1', + config=config, + data=data, + metadata=metadata['metadata'], + ) + assert output_data is False + assert opc_metrics == {'tag1': side_effect[0], 'tag2': side_effect[1]} + opc.write_data.assert_has_calls( + [ + call( + server_id='server1', + tag='tag1', + data=0.75, + data_type='float', + tag_type='prediction', + metadata=metadata['metadata'], + ), + call( + server_id='server1', + tag='tag2', + data=0.95, + data_type='float', + tag_type='confidence', + metadata=metadata['metadata'], + ), + ] + ) + + +@mark.asyncio +async def test_manage_output_tags_do_nothing(opc): + opc.write_data = AsyncMock(return_value=0.1) + data = DataFrame({'prediction': [0.75], 'prediction_confidence': [0.95]}) + config = { + '_invalid_key': {'tag1': {'data_type': 'float'}}, + } + output_data, opc_metrics = await opc.manage_output_tags( + server_id='server1', + config=config, + data=data, + metadata=metadata['metadata'], + ) + assert output_data is True + assert opc_metrics == {} + opc.write_data.assert_not_called() + + +@mark.asyncio +@patch('laborious.activities.opc.DataFrame') +async def test_write_opc_data_success(mock_dataframe, opc): # Arrange input_data = { **metadata, @@ -270,37 +371,25 @@ 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) + output_data, opc_metrics = 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'} + assert opc_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'], ) - 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..4da51ee 100644 --- a/tests/laborious/utils/repository/test_opc_repository.py +++ b/tests/laborious/utils/repository/test_opc_repository.py @@ -1,6 +1,6 @@ import json from datetime import datetime -from unittest.mock import ANY, AsyncMock, MagicMock, Mock, call, patch +from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch import pytest from asyncua.crypto.security_policies import SecurityPolicyBasic256 @@ -236,22 +236,22 @@ async def test_validate_connection_none_client(opc_repository): opc_repository.connect.assert_called_once() -@pytest.mark.asyncio -async def test_validate_connection_error_count_disconnect_error(opc_repository): - opc_repository.error_count = 6 - opc_repository.client = AsyncMock() - opc_repository.disconnect = AsyncMock(side_effect=Exception('Test error')) - opc_repository.connect = AsyncMock(return_value=(True, {})) +# @pytest.mark.asyncio +# async def test_validate_connection_error_count_disconnect_error(opc_repository): +# opc_repository.error_count = 6 +# opc_repository.client = AsyncMock() +# opc_repository.disconnect = AsyncMock(side_effect=Exception('Test error')) +# opc_repository.connect = AsyncMock(return_value=(True, {})) - response = await opc_repository.validate_connection() - assert response == opc_repository.connect.return_value - opc_repository.disconnect.assert_called_once() - opc_repository.connect.assert_called_once() - opc_repository.logger.custom_error.assert_has_calls( - [ - call('Failed to disconnect from OPC server: Test error', ANY), - ] - ) +# response = await opc_repository.validate_connection() +# assert response == opc_repository.connect.return_value +# opc_repository.disconnect.assert_called_once() +# opc_repository.connect.assert_called_once() +# opc_repository.logger.custom_error.assert_has_calls( +# [ +# call('Failed to disconnect from OPC server: Test error', ANY), +# ] +# ) @pytest.mark.asyncio @@ -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..ecd3bc2 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,12 +93,27 @@ 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, - **metadata, + 'data': prediction_data, 'timestamp_conversion': { 'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ, }, + **metadata, + }, + retry_policy=ANY, + start_to_close_timeout=ANY, + ) + ] + ) + + 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, @@ -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 diff --git a/tests/laborious/workflows/subworkflows/test_prediction_process.py b/tests/laborious/workflows/subworkflows/test_prediction_process.py index 6baec30..0d8a0c1 100644 --- a/tests/laborious/workflows/subworkflows/test_prediction_process.py +++ b/tests/laborious/workflows/subworkflows/test_prediction_process.py @@ -170,7 +170,7 @@ async def test_run(workflow_mock, prediction_process): ) workflow_mock.execute_child_workflow.assert_called_once_with( - 'format_and_export_prediction', + 'subworkflow.format_and_export_prediction', { 'metadata': metadata, 'path_flag': 'continue', @@ -730,7 +730,7 @@ async def test_path_flag_handler_continue(workflow_mock, prediction_process): assert result is True workflow_mock.execute_activity_method.assert_not_called() workflow_mock.execute_child_workflow.assert_called_once_with( - 'format_and_export_prediction', + 'subworkflow.format_and_export_prediction', { 'metadata': metadata, 'path_flag': path_flag, diff --git a/tests/laborious/workflows/test_predictions_batch.py b/tests/laborious/workflows/test_predictions_batch.py index 62cc43c..7979c47 100644 --- a/tests/laborious/workflows/test_predictions_batch.py +++ b/tests/laborious/workflows/test_predictions_batch.py @@ -75,5 +75,5 @@ async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch } workflow_mock.execute_child_workflow.assert_has_calls( - [call('prediction_process', prediction_input)] + [call('subworkflow.prediction_process', prediction_input)] ) diff --git a/values.yaml b/values.yaml index e4b3165..d5faf9b 100644 --- a/values.yaml +++ b/values.yaml @@ -151,7 +151,7 @@ env: - name: GITHUB_REPO_URL value: "git@github.com:Aignosi/sientia-dataops-laborious_temporal.git" - name: GITHUB_BRANCH - value: "SIENTIAPDE-1312-melhorias-e-correcoes-nas-pipelines-de-dados" + value: "fix/SIENTIAPDE-1314-ajustes-nas-camadas-de-monitoramento-do-sientia" - name: PYTHON_APP value: "laborious.worker.worker"