SIENTIAPDE-1325

Update logging in MLFlowRepository and OpcRepository to use unified logging methods

- Refactored logging calls in MLFlowRepository to replace `self.logger.info` and `self.logger.debug` with `self.info` and `self.debug` for consistency.
- Updated connection logging in OpcRepository to format the connection message properly.
- Adjusted test cases to reflect changes in logging behavior and ensure proper assertions.

.github/workflows/release.yml

@@ -8,6 +8,7 @@ on:
 
 jobs:
   release:
+    if: github.event.pull_request.merged == true
     uses: Aignosi/github_workflow_templates/.github/workflows/dataops-module-release.yml@main
     permissions: write-all
     with:

laborious/activities/opc.py

@@ -105,7 +105,9 @@ class OPC(SientiaMonitoring):
                     attachment_content=error_data.get('attachment_content', None),
                 )
             else:
-                self.logger.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.')
+                self.logger.info(
+                    f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
+                )
 
     async def write_data(
         self,

laborious/utils/repository/model_repository.py

@@ -243,7 +243,7 @@ class MLFlowRepository(SientiaMonitoring):
             rmtree(full_path)
         makedirs(output_dir, exist_ok=True)
 
-        self.logger.info(f'Downloading artifacts from {run_id} to {output_dir}')
+        self.info(f'Downloading artifacts from {run_id} to {output_dir}')
 
         core_labels = self.get_core_labels(metadata, operation_type='download_artifacts')
 
@@ -279,7 +279,7 @@ class MLFlowRepository(SientiaMonitoring):
             - Warnings during the model loading process are suppressed.
         """
         model_uri = f'models:/{model_name}/production'
-        self.logger.info(f'Loading prediction model {model_name} from {model_uri}')
+        self.info(f'Loading prediction model {model_name} from {model_uri}')
 
         core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
         start_time = time.time()
@@ -326,7 +326,7 @@ class MLFlowRepository(SientiaMonitoring):
         latest_production_id = self.get_model_run_id(model_name=model_name, stage='Production')
         model_uri = self.get_model_uri(latest_production_id, prediction=False)
 
-        self.logger.info(f'Loading data model {model_name} from {model_uri}')
+        self.info(f'Loading data model {model_name} from {model_uri}')
 
         core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
         start_time = time.time()
@@ -370,7 +370,7 @@ class MLFlowRepository(SientiaMonitoring):
             tuple[Any, str | None]: Model object and optional artifact path.
         """
 
-        self.logger.info(
+        self.info(
             f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}'
         )
 
@@ -380,15 +380,13 @@ class MLFlowRepository(SientiaMonitoring):
         artifact_path = None
 
         if load_wrapper:
-            self.logger.info(
+            self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
-                f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}'
-            )
 
             target = 'prediction_model' if model_type == 'predict' else 'data_model'
 
             artifact_path = await self.dowload_artifacts(model_name, metadata, target)
 
-            self.logger.info(
+            self.info(
                 f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
             )
 
@@ -423,12 +421,16 @@ class MLFlowRepository(SientiaMonitoring):
         Returns:
             pd.DataFrame: DataFrame with converted datetime index.
         """
+        if data.empty:
+            self.info('Data is empty, skipping datetime index detection and parsing', metadata)
+            return data
+
         index = data.index
 
         # Get type of first element of index
         index_type = type(index[0])
 
-        self.logger.custom_info(f'Index type: {index_type}', metadata)
+        self.info(f'Index type: {index_type}', metadata)
 
         message = f'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format {DATETIME_FORMAT_WITH_TZ}.'
 
@@ -488,7 +490,7 @@ class MLFlowRepository(SientiaMonitoring):
         Returns:
             dict: Model configuration.
         """
-        self.logger.debug(f'Model {model_name} is still valid, using cached version')
+        self.debug(f'Model {model_name} is still valid, using cached version')
 
         return cache['target']
 
@@ -503,7 +505,7 @@ class MLFlowRepository(SientiaMonitoring):
         Returns:
             None
         """
-        self.logger.debug(f'Model {model_name} is outdated, downloading a new one')
+        self.debug(f'Model {model_name} is outdated, downloading a new one')
 
         del self.model_cache[model_key]['target']
         del self.model_cache[model_key]
@@ -554,7 +556,7 @@ class MLFlowRepository(SientiaMonitoring):
                     # Model is outdated, delete old model files
                     self.handle_outdated_model(model_name=model_name, model_key=model_key)
             else:
-                self.logger.debug(
+                self.debug(
                     f'Model {model_name} is not in {model_type} cache, downloading a new one'
                 )
 
@@ -631,7 +633,7 @@ class MLFlowRepository(SientiaMonitoring):
         prediction = model.predict(data)
 
         if retention == 0:
-            self.logger.info(f'Deleting model {model_name}:{operation} from memory')
+            self.info(f'Deleting model {model_name}:{operation} from memory')
             del model
 
             force_memory_release(self.logger)
@@ -675,8 +677,8 @@ class MLFlowRepository(SientiaMonitoring):
                 `data_model`, including optional artifact paths.
         """
 
-        self.logger.custom_info(f'Starting model experiment creation for {model_name}', metadata)
+        self.info(f'Starting model experiment creation for {model_name}', metadata)
-        self.logger.custom_debug(
+        self.debug(
             f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}',
             metadata,
         )
@@ -684,10 +686,8 @@ class MLFlowRepository(SientiaMonitoring):
         # data.to_csv(
         #     f"tmp/retrain_data_{model_name}.csv", index=True)
 
-        self.logger.custom_info(
+        self.info(f'Retrieved latest production run ID: {latest_production_id}', metadata)
-            f'Retrieved latest production run ID: {latest_production_id}', metadata
+        self.info(f'Loading transformation model for {model_name}', metadata)
-        )
-        self.logger.custom_info(f'Loading transformation model for {model_name}', metadata)
 
         load_transform_wrapper = transform_flavor == 'pyfunc'
 
@@ -699,7 +699,7 @@ class MLFlowRepository(SientiaMonitoring):
             load_wrapper=load_transform_wrapper,
         )
 
-        self.logger.custom_info(f'Loading prediction model for {model_name}', metadata)
+        self.info(f'Loading prediction model for {model_name}', metadata)
 
         load_predict_wrapper = predict_flavor == 'pyfunc'
 
@@ -728,24 +728,22 @@ class MLFlowRepository(SientiaMonitoring):
 
         treated_data = treated_data.drop_duplicates(subset=['timestamp'], keep='first')
 
-        self.logger.custom_debug(f'Treated data index: {treated_data.index}', metadata)
+        self.debug(f'Treated data index: {treated_data.index}', metadata)
 
         # treated_data.to_csv(
         #     f"tmp/retrain_treated_data_{model_name}.csv", index=True)
 
-        self.logger.custom_debug(f'Transformed data shape: {treated_data.shape}', metadata)
+        self.debug(f'Transformed data shape: {treated_data.shape}', metadata)
 
         if target_name is None:
             target_name = data_model.target_variable
-            self.logger.custom_debug(
+            self.debug(f'Using target variable from data model: {target_name}', metadata)
-                f'Using target variable from data model: {target_name}', metadata
-            )
         else:
-            self.logger.custom_debug(f'Using provided target variable: {target_name}', metadata)
+            self.debug(f'Using provided target variable: {target_name}', metadata)
 
         # Check if treated_data contains target variable
         if target_name not in treated_data.columns:
-            self.logger.custom_debug(
+            self.debug(
                 f'Target variable {target_name} not found in treated data, aligning data with treated data indexes',
                 metadata,
             )
@@ -757,9 +755,7 @@ class MLFlowRepository(SientiaMonitoring):
             )
         else:
             # Uses target variable from treated data
-            self.logger.custom_debug(
+            self.debug(f'Target variable {target_name} found in treated data, using it', metadata)
-                f'Target variable {target_name} found in treated data, using it', metadata
-            )
             retrain_dataset = treated_data
 
         # retrain_dataset.to_csv(
@@ -767,9 +763,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         prediction_model.fit(retrain_dataset)
 
-        self.logger.custom_info(
+        self.info(f'Model experiment creation completed successfully for {model_name}', metadata)
-            f'Model experiment creation completed successfully for {model_name}', metadata
-        )
 
         retrain_data = {
             'prediction_model': {
@@ -791,7 +785,7 @@ class MLFlowRepository(SientiaMonitoring):
         """
         model = model_data['model']
 
-        self.logger.custom_debug(f'Logging {model_type} model to {model_type}', metadata)
+        self.debug(f'Logging {model_type} model to {model_type}', metadata)
 
         core_labels = self.get_core_labels(metadata, operation_type='log_model')
         start_time = time.time()
@@ -802,11 +796,11 @@ class MLFlowRepository(SientiaMonitoring):
             elif flavor == 'pyfunc':
                 code_path = [path.join(model_data['artifact_path'], 'code', 'utils')]
 
-                self.logger.custom_debug(f'Code path: {code_path}', metadata)
+                self.debug(f'Code path: {code_path}', metadata)
 
                 model.store_model(artifact_path=model_type, code_path=code_path, to_disk=False)
 
-                self.logger.custom_debug('Model uploaded successfully', metadata)
+                self.debug('Model uploaded successfully', metadata)
             elif flavor == 'pytorch':
                 mlflow.pytorch.log_model(model, model_type)
             else:
@@ -858,7 +852,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         model_temp_path = path.join(ARTIFACTS_PATH, model_name)
 
-        self.logger.custom_info(f'Starting model retraining process for {model_name}', metadata)
+        self.info(f'Starting model retraining process for {model_name}', metadata)
 
         original_params = self.get_model_params(latest_production_id)
         retrain_params = {
@@ -875,7 +869,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         current_run_name = self.get_next_run_name(experiment_name)
 
-        self.logger.custom_debug(f'Attributes: {retrain_params}', metadata)
+        self.debug(f'Attributes: {retrain_params}', metadata)
 
         data_path = f'{model_temp_path}/retrain_data.csv'
 
@@ -883,7 +877,7 @@ class MLFlowRepository(SientiaMonitoring):
 
         data.to_csv(data_path, index=True)
 
-        self.logger.custom_info(
+        self.info(
             f'Starting model upload for {experiment_name} with run name {current_run_name}',
             metadata,
         )
@@ -897,17 +891,17 @@ class MLFlowRepository(SientiaMonitoring):
                 description=experiment_description,
             ) as _run:
                 run_id = _run.info.run_id
-                self.logger.custom_info('Logging data model', metadata)
+                self.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.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.info(f'Model logged successfully for {model_name}', metadata)
 
-                self.logger.custom_info(f'Logging remaining parameters for {model_name}', metadata)
+                self.info(f'Logging remaining parameters for {model_name}', metadata)
 
                 # update transfomation model
                 # fixed parameters
@@ -923,15 +917,15 @@ class MLFlowRepository(SientiaMonitoring):
         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)
+        self.info('Deleting model from filesystem', metadata)
         if path.exists(model_temp_path):
             rmtree(model_temp_path)
 
-        self.logger.custom_info('Deleting prediction model from memory', metadata)
+        self.info('Deleting prediction model from memory', metadata)
         del prediction_model['model']
         del prediction_model
 
-        self.logger.custom_info('Deleting data model from memory', metadata)
+        self.info('Deleting data model from memory', metadata)
         del data_model['model']
         del data_model
 
@@ -971,7 +965,7 @@ class MLFlowRepository(SientiaMonitoring):
             4. Archives existing production versions
         """
 
-        self.logger.custom_info(
+        self.info(
             f'Starting production model update for {model_name} with run ID: {run_id}', metadata
         )
 
@@ -1059,9 +1053,7 @@ class MLFlowRepository(SientiaMonitoring):
                       and returned in the response structure rather than propagated.
         """
 
-        self.logger.custom_debug(
+        self.debug(f'Data received for model transformation: {data.head(5).to_csv()}', metadata)
-            f'Data received for model transformation: {data.head(5).to_csv()}', metadata
-        )
 
         # data.to_csv(
         #     f"tmp/data_{model_name}.csv", index=True)
@@ -1079,7 +1071,7 @@ class MLFlowRepository(SientiaMonitoring):
                 metadata=metadata,
             )
 
-            self.logger.custom_debug(
+            self.debug(
                 f'Data received from model transformation: {transformed_data.head(5).to_csv()}',
                 metadata,
             )
@@ -1147,9 +1139,7 @@ class MLFlowRepository(SientiaMonitoring):
             input_index = data.index
             start_time = datetime.now()
 
-            self.logger.custom_debug(
+            self.debug(f'Data received for model prediction: {data.head(5).to_csv()}', metadata)
-                f'Data received for model prediction: {data.head(5).to_csv()}', metadata
-            )
 
             # data.to_csv(
             #     f"tmp/treated_data_{model_name}.csv", index=True)
@@ -1166,7 +1156,7 @@ class MLFlowRepository(SientiaMonitoring):
             end_time = datetime.now()
 
             if isinstance(predict_data, pd.DataFrame):
-                self.logger.custom_debug(
+                self.debug(
                     f'Data received from model prediction: {data.head(5).to_csv()}', metadata
                 )
 
@@ -1235,22 +1225,22 @@ class MLFlowRepository(SientiaMonitoring):
             Exception: Any other exception during the retraining process
         """
 
-        self.logger.custom_info(f'Starting model retraining workflow for {model_name}', metadata)
+        self.info(f'Starting model retraining workflow for {model_name}', metadata)
-        self.logger.custom_debug(f'Data received for model retraining: {data.to_csv()}', metadata)
+        self.debug(f'Data received for model retraining: {data.to_csv()}', metadata)
 
         target_name = model_config.get('target', None)
 
         transform_flavor = model_config.get('transform_flavor', 'sklearn')
         predict_flavor = model_config.get('predict_flavor', 'sklearn')
 
-        self.logger.custom_debug(
+        self.debug(
             f'Model configuration - transform_flavor: {transform_flavor}, predict_flavor: {predict_flavor}, target_name: {target_name}',
             metadata,
         )
 
         try:
             latest_production_id = self.get_model_run_id(model_name, stage='Production')
-            self.logger.custom_info('Creating model experiment environment', metadata)
+            self.info('Creating model experiment environment', metadata)
             retrain_data = await self.fit_models(
                 model_name=model_name,
                 data=data,
@@ -1260,11 +1250,9 @@ class MLFlowRepository(SientiaMonitoring):
                 metadata=metadata,
                 latest_production_id=latest_production_id,
             )
-            self.logger.custom_info(
+            self.info(f'Model experiment created successfully: {retrain_data}', metadata)
-                f'Model experiment created successfully: {retrain_data}', metadata
-            )
 
-            self.logger.custom_info('Saving model retrain', metadata)
+            self.info('Saving model retrain', metadata)
             experiment = await self.create_new_experiment(
                 model_name=model_name,
                 data=data,
@@ -1274,7 +1262,7 @@ class MLFlowRepository(SientiaMonitoring):
                 metadata=metadata,
                 latest_production_id=latest_production_id,
             )
-            self.logger.custom_info(
+            self.info(
                 f'Model retraining completed successfully for experiment: {experiment}', metadata
             )
 
@@ -1285,7 +1273,7 @@ class MLFlowRepository(SientiaMonitoring):
             }
         except Exception as e:
             error_msg = f'Error retraining model {model_name}: {e}'
-            self.logger.custom_info(error_msg, metadata)
+            self.info(error_msg, metadata)
             return {
                 'success': False,
                 'experiment': None,

laborious/utils/repository/opc_repository.py

@@ -143,7 +143,9 @@ class OpcRepository(SientiaMonitoring):
 
         if self.cert_path:
             await self.set_security()
-        self.logger.custom_info(f'Starting connection to OPC server {self.id}:{self.server_name}...', self.metadata)
+        self.logger.custom_info(
+            f'Starting connection to OPC server {self.id}:{self.server_name}...', self.metadata
+        )
         return await self.try_connect()
 
     async def try_connect(self) -> tuple[bool, dict[str, Any]]:
@@ -176,7 +178,7 @@ class OpcRepository(SientiaMonitoring):
                     'level': NotificationLevel.ERROR,
                 }
             await self.client.connect()
-        
+
             await self.emit_metric(
                 metric_object=metrics.OPC_CONNECTION_STATUS,
                 method='set',

tests/laborious/activities/test_mlflow.py

@@ -255,9 +255,7 @@ async def test_retrain_model_success_data_success_retrain(mock_to_datetime, mlfl
     timestamp = raw_data.__getitem__.return_value.max.return_value
 
     raw_data.sort_values.assert_not_called()
-    raw_data.drop_duplicates.assert_called_once_with(
+    raw_data.drop_duplicates.assert_called_once_with(subset=['variable', 'timestamp'], keep='first')
-        subset=['variable', 'timestamp'], keep='first'
-    )
     raw_data = raw_data.drop_duplicates.return_value
 
     raw_data.drop.assert_has_calls(

tests/laborious/activities/test_opc.py

@@ -61,6 +61,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
     mock_notification_handler = MagicMock()
     servers = {
         'server1': {
+            'server_name': 'server1',
             'id': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -70,6 +71,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'reconnection_interval': 60,
         },
         'server2': {
+            'server_name': 'server2',
             'id': 'server2',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -79,6 +81,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
             'reconnection_interval': 60,
         },
         'server3': {
+            'server_name': 'server3',
             'id': 'server3',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
@@ -106,6 +109,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
         [
             call(
                 opc_id='server1',
+                server_name='server1',
                 url='http://localhost:8080',
                 logger=mock_logger,
                 server_uri='opc.tcp://localhost:4840',
@@ -122,6 +126,7 @@ async def test_init_opc(mock_send_notification, mock_opc_repository):
         [
             call(
                 opc_id='server2',
+                server_name='server2',
                 url='http://localhost:8080',
                 logger=mock_logger,
                 server_uri='opc.tcp://localhost:4840',
@@ -163,6 +168,7 @@ async def opc(mock_opc_repository):
     servers = {
         'server1': {
             'id': 'server1',
+            'server_name': 'server1',
             'url': 'http://localhost:8080',
             'server_uri': 'opc.tcp://localhost:4840',
             'cert_path': '',

tests/laborious/utils/repository/test_minio_repository.py

@@ -95,7 +95,6 @@ async def test_create_bucket_error(minio_repository):
     with raises(ValueError):
         await minio_repository.create_bucket({})
 
-
         minio_repository.s3_client.create_bucket.assert_called_once_with(Bucket='test')
         minio_repository.emit_metric.assert_called_once_with(
             metric_object=metrics.MINIO_WRITE_ERROR_COUNT, tags=ANY

tests/laborious/utils/repository/test_opc_repository.py

@@ -18,6 +18,7 @@ def mock_logger():
 def opc_repository(mock_logger):
     repository = OpcRepository(
         opc_id='test_repo',
+        server_name='test_server',
         url='opc.tcp://localhost:4840',
         logger=mock_logger,
         notification_handler=Mock(),
@@ -55,6 +56,7 @@ metadata = {
 
 def test_init(opc_repository):
     assert opc_repository.id == 'test_repo'
+    assert opc_repository.server_name == 'test_server'
     assert opc_repository.url == 'opc.tcp://localhost:4840'
     assert opc_repository.server_uri == 'urn:test:server'
     assert opc_repository.cert_path == '/path/to/cert.pem'
@@ -139,11 +141,13 @@ async def test_try_connect_success(opc_repository):
 @pytest.mark.asyncio
 async def test_try_connect_fail(opc_repository):
     opc_repository.last_reconnection_time = None
+    opc_repository.disconnect = AsyncMock()
     opc_repository.client = MagicMock()
     opc_repository.client.connect.side_effect = Exception('Test error')
 
     is_connected, error_data = await opc_repository.try_connect()
 
+    opc_repository.disconnect.assert_called_once()
     opc_repository.client.connect.assert_called_once()
     assert is_connected is False
     assert error_data['notification_id'] == f'OPC_CONNECTION_ERROR_{opc_repository.id}'