QTZPOC-20

QTZPOC-20: Implement structured logging with metadata for repository, worker, and OPC activities.

laborious/activities/opc.py

@@ -122,7 +122,7 @@ class OPC(SientiaMonitoring):
             misconfigured server intentionally stops every other server, and the
             whole worker, from starting.
         """
-        self.logger.info('Initializing OPC servers...')
+        self.info('Initializing OPC servers...')
         for opc_id, server in self.opc_servers.items():
             self.opc_repository[opc_id] = OpcRepository(
                 opc_id=server['id'],
@@ -153,9 +153,7 @@ class OPC(SientiaMonitoring):
                     attachment_content=error_data.get('attachment_content', None),
                 )
             else:
-                self.logger.info(
+                self.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.')
-                    f'OPC server {opc_id}:{server["server_name"]} connected successfully.'
-                )
 
     async def write_data(
         self,

laborious/utils/repository/model_repository.py

@@ -52,7 +52,7 @@ MODEL_PICKLE_NAME = 'model.pkl'
 INVALID_FLAVOR_MESSAGE = "Invalid flavor. Use 'sklearn' or 'pyfunc' or 'joblib'."
 
 
-def force_memory_release(logger: Logger):
+def force_memory_release(logger: Logger, metadata: dict[str, Any] | None = None):
     """Attempt to release memory from the Python process.
 
     Executes a garbage collection cycle and calls `malloc_trim(0)` on glibc
@@ -61,14 +61,17 @@ def force_memory_release(logger: Logger):
 
     Args:
         logger (Logger): Logger for observability.
+        metadata (dict[str, Any] | None): Metadata used for structured logging.
     """
     gc.collect()
 
+    metadata = metadata or {}
+
     try:
         ctypes.CDLL('libc.so.6').malloc_trim(0)
-        logger.info('Memory released')
+        logger.custom_info('Memory released', metadata)
     except Exception as e:
-        logger.info(f'Memory release failed: {e}')
+        logger.custom_info(f'Memory release failed: {e}', metadata)
 
 
 class MLFlowRepository(SientiaMonitoring):
@@ -320,7 +323,7 @@ class MLFlowRepository(SientiaMonitoring):
             rmtree(full_path)
         makedirs(output_dir, exist_ok=True)
 
-        self.info(f'Downloading artifacts from {run_id} to {output_dir}')
+        self.info(f'Downloading artifacts from {run_id} to {output_dir}', metadata)
 
         core_labels = self.get_core_labels(metadata, operation_type='download_artifacts')
 
@@ -676,7 +679,7 @@ class MLFlowRepository(SientiaMonitoring):
         model_uri = None if flavor == 'joblib' else f'models:/{model_name}/production'
 
         if model_uri is not None:
-            self.info(f'Loading prediction model {model_name} from {model_uri}')
+            self.info(f'Loading prediction model {model_name} from {model_uri}', metadata)
 
         core_labels = self.get_core_labels(metadata, operation_type='load_predict_model')
         start_time = time.time()
@@ -756,7 +759,7 @@ class MLFlowRepository(SientiaMonitoring):
         )
 
         if model_uri is not None:
-            self.info(f'Loading data model {model_name} from {model_uri}')
+            self.info(f'Loading data model {model_name} from {model_uri}', metadata)
 
         core_labels = self.get_core_labels(metadata, operation_type='load_transform_model')
         start_time = time.time()
@@ -822,7 +825,8 @@ class MLFlowRepository(SientiaMonitoring):
         """
 
         self.info(
-            f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}'
+            f'Downloading {model_type} model {model_name} with flavor {flavor} and load_wrapper {load_wrapper}',
+            metadata,
         )
 
         if model_type not in ['predict', 'transform']:
@@ -838,7 +842,10 @@ class MLFlowRepository(SientiaMonitoring):
 
         if load_wrapper:
             core_labels = self.get_core_labels(metadata, operation_type='download_model_wrapper')
-            self.info(f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}')
+            self.info(
+                f'Loading wrapper for {model_type} model {model_name} with flavor {flavor}',
+                metadata,
+            )
 
             try:
                 artifact_path = await self.dowload_artifacts(
@@ -846,7 +853,8 @@ class MLFlowRepository(SientiaMonitoring):
                 )
 
                 self.info(
-                    f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}'
+                    f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}',
+                    metadata,
                 )
 
                 raw_model = mlflow.pyfunc.load_model(artifact_path)
@@ -959,33 +967,37 @@ class MLFlowRepository(SientiaMonitoring):
             return False
         return True
 
-    def handle_valid_model(self, model_name: str, cache: dict) -> dict:
+    def handle_valid_model(self, model_name: str, cache: dict, metadata: dict[str, Any]) -> dict:
         """
         Return the cached model configuration when retention is valid.
 
         Args:
             model_name (str): Name of the model (for logging/consistency).
             cache (dict): Cached model data structure.
+            metadata (dict[str, Any]): Metadata used for structured logging.
 
         Returns:
             dict: Model configuration.
         """
-        self.debug(f'Model {model_name} is still valid, using cached version')
+        self.debug(f'Model {model_name} is still valid, using cached version', metadata)
 
         return cache['target']
 
-    def handle_outdated_model(self, model_name: str, model_key: str) -> None:
+    def handle_outdated_model(
+        self, model_name: str, model_key: str, metadata: dict[str, Any]
+    ) -> None:
         """
         Clean up outdated cached model and its artifacts.
 
         Args:
             model_name (str): Name of the model.
             model_key (str): Cache key for the model.
+            metadata (dict[str, Any]): Metadata used for structured logging.
 
         Returns:
             None
         """
-        self.debug(f'Model {model_name} is outdated, downloading a new one')
+        self.debug(f'Model {model_name} is outdated, downloading a new one', metadata)
 
         del self.model_cache[model_key]['target']
         del self.model_cache[model_key]
@@ -1031,13 +1043,18 @@ class MLFlowRepository(SientiaMonitoring):
 
                 # 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)
+                    return self.handle_valid_model(
+                        model_name=model_name, cache=cache, metadata=metadata
+                    )
                 else:
                     # Model is outdated, delete old model files
-                    self.handle_outdated_model(model_name=model_name, model_key=model_key)
+                    self.handle_outdated_model(
+                        model_name=model_name, model_key=model_key, metadata=metadata
+                    )
             else:
                 self.debug(
-                    f'Model {model_name} is not in {model_type} cache, downloading a new one'
+                    f'Model {model_name} is not in {model_type} cache, downloading a new one',
+                    metadata,
                 )
 
         # Donwload new model (without lock to avoid blocking other threads)
@@ -1113,10 +1130,10 @@ class MLFlowRepository(SientiaMonitoring):
         prediction = model.predict(data)
 
         if retention == 0:
-            self.info(f'Deleting model {model_name}:{operation} from memory')
+            self.info(f'Deleting model {model_name}:{operation} from memory', metadata)
             del model
 
-            force_memory_release(self.logger)
+            force_memory_release(self.logger, metadata)
 
         return prediction
 
@@ -1550,7 +1567,7 @@ class MLFlowRepository(SientiaMonitoring):
         del data_model['model']
         del data_model
 
-        force_memory_release(self.logger)
+        force_memory_release(self.logger, metadata)
 
         return {
             'run_id': run_id,

laborious/worker/worker.py

@@ -136,7 +136,7 @@ async def main():
     metadata_runtime = {**metadata, 'runtime': runtime}
 
     logger.custom_info('Starting prometheus client...', metadata_runtime)
-    start_prometheus_server()
+    start_prometheus_server(logger, metadata_runtime)
 
     logger.custom_info('Starting Notification Handler...', metadata_runtime)
 
@@ -313,7 +313,7 @@ async def main():
         exit_code = e.code if isinstance(e, SystemExit) and isinstance(e.code, int) else 0
         raise
     except Exception as e:
-        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
+        logger.custom_error(f'An unhandled exception occurred: {e}', metadata_runtime)
         exit_code = 1
     finally:
         if notification_handler:
@@ -324,7 +324,7 @@ async def main():
         sys.exit(exit_code)
 
 
-def start_prometheus_server():
+def start_prometheus_server(logger, metadata: dict):
     """
     Starts the Prometheus metrics server for monitoring and observability.
 
@@ -344,10 +344,10 @@ def start_prometheus_server():
     try:
         port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f'Prometheus server started on port {port}.')
+        logger.custom_info(f'Prometheus server started on port {port}.', metadata)
         metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
     except Exception as e:
-        print(f'Failed to start Prometheus server: {e}')
+        logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
         os._exit(1)
 
 

tests/laborious/utils/repository/test_model_repository.py

@@ -27,25 +27,27 @@ from laborious.utils.repository.model_repository import (
 @patch('laborious.utils.repository.model_repository.gc')
 def test_force_memory_release_success(gc, ctypes):
     logger = MagicMock()
+    metadata = {'model_name': 'my-model'}
 
-    force_memory_release(logger)
+    force_memory_release(logger, metadata)
 
     gc.collect.assert_called_once()
     ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0)
-    logger.info.assert_called_once_with('Memory released')
+    logger.custom_info.assert_called_once_with('Memory released', metadata)
 
 
 @patch('laborious.utils.repository.model_repository.ctypes')
 @patch('laborious.utils.repository.model_repository.gc')
 def test_force_memory_release_error(gc, ctypes):
     logger = MagicMock()
+    metadata = {'model_name': 'my-model'}
     ctypes.CDLL.return_value.malloc_trim.side_effect = Exception('error')
-    force_memory_release(logger)
+    force_memory_release(logger, metadata)
 
     gc.collect.assert_called_once()
     ctypes.CDLL.return_value.malloc_trim.assert_called_once_with(0)
 
-    logger.info.assert_called_once_with('Memory release failed: error')
+    logger.custom_info.assert_called_once_with('Memory release failed: error', metadata)
 
 
 @pytest.fixture
@@ -807,7 +809,7 @@ def test_check_cache_retention_true(datetime_mock, mlflow_repository):
 
 def test_handle_valid_model(mlflow_repository):
     cache = {'target': {'model': 'model', 'artifact_path': 'test_artifact_path'}}
-    output = mlflow_repository.handle_valid_model('model_name', cache)
+    output = mlflow_repository.handle_valid_model('model_name', cache, {})
     assert output == {'model': 'model', 'artifact_path': 'test_artifact_path'}
 
 
@@ -817,7 +819,7 @@ def test_handle_outdated_model(mlflow_repository):
             'target': {'model': 'model', 'artifact_path': 'test_artifact_path'}
         }
     }
-    mlflow_repository.handle_outdated_model('model_name', 'model_name_transform')
+    mlflow_repository.handle_outdated_model('model_name', 'model_name_transform', {})
     assert mlflow_repository.model_cache == {}
 
 
@@ -857,7 +859,9 @@ async def test_get_model_cached_valid(mlflow_repository):
     )
 
     mlflow_repository.handle_valid_model.assert_called_once_with(
-        model_name='model_name', cache=mlflow_repository.model_cache['model_name_predict']
+        model_name='model_name',
+        cache=mlflow_repository.model_cache['model_name_predict'],
+        metadata={},
     )
 
     mlflow_repository.handle_outdated_model.assert_not_called()
@@ -887,7 +891,7 @@ async def test_get_model_cached_outdated(mlflow_repository):
     )
     mlflow_repository.handle_valid_model.assert_not_called()
     mlflow_repository.handle_outdated_model.assert_called_once_with(
-        model_name='model_name', model_key='model_name_predict'
+        model_name='model_name', model_key='model_name_predict', metadata={}
     )
 
 
@@ -916,7 +920,7 @@ async def test_get_cached_operation_retention_0(force_memory_release, mlflow_rep
         'model_name', data, 'transform', 0, 'sklearn', {}
     )
     assert output == model.predict.return_value
-    force_memory_release.assert_called_once_with(mlflow_repository.logger)
+    force_memory_release.assert_called_once_with(mlflow_repository.logger, {})
 
 
 @patch('laborious.utils.repository.model_repository.force_memory_release')
@@ -1316,7 +1320,7 @@ async def test_create_new_experiment(
         ]
     )
 
-    force_memory_release.assert_called_once_with(mlflow_repository.logger)
+    force_memory_release.assert_called_once_with(mlflow_repository.logger, metadata['metadata'])
 
     assert report == {
         'run_id': mlflow.start_run.return_value.__enter__.return_value.info.run_id,

tests/laborious/worker/test_worker_main.py

@@ -210,7 +210,7 @@ class TestMainHappyPath:
 @patch('laborious.worker.worker.metrics.APP_UP')
 @patch('laborious.worker.worker.start_http_server')
 def test_start_prometheus_server_marks_the_app_up(mock_start_http_server, mock_app_up):
-    start_prometheus_server()
+    start_prometheus_server(MagicMock(), {})
 
     mock_start_http_server.assert_called_once()
     mock_app_up.labels.return_value.set.assert_called_once_with(1)
@@ -224,7 +224,7 @@ def test_start_prometheus_server_exits_the_process_on_failure(
 ):
     mock_start_http_server.side_effect = OSError('port in use')
 
-    start_prometheus_server()
+    start_prometheus_server(MagicMock(), {})
 
     mock_os_exit.assert_called_once_with(1)
     mock_app_up.labels.return_value.set.assert_not_called()
@@ -238,6 +238,6 @@ def test_start_prometheus_server_reads_the_configured_port(
 ):
     monkeypatch.setenv('HTTP_METRICS_PORT', port_env)
 
-    start_prometheus_server()
+    start_prometheus_server(MagicMock(), {})
 
     mock_start_http_server.assert_called_once_with(int(port_env))