diff --git a/laborious/activities/opc.py b/laborious/activities/opc.py index 2c14cf5..7dac8da 100644 --- a/laborious/activities/opc.py +++ b/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( - f'OPC server {opc_id}:{server["server_name"]} connected successfully.' - ) + self.info(f'OPC server {opc_id}:{server["server_name"]} connected successfully.') async def write_data( self, diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 8d3d9f6..285957a 100644 --- a/laborious/utils/repository/model_repository.py +++ b/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, diff --git a/laborious/worker/worker.py b/laborious/worker/worker.py index e585a2d..d096820 100644 --- a/laborious/worker/worker.py +++ b/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) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 0311ef1..454605a 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/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, diff --git a/tests/laborious/worker/test_worker_main.py b/tests/laborious/worker/test_worker_main.py index 92e872f..a222464 100644 --- a/tests/laborious/worker/test_worker_main.py +++ b/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))