From c12f2996c2aaa8e998b3be2b0d425b5c77696da0 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Tue, 11 Aug 2026 11:36:22 -0300 Subject: [PATCH 1/4] QTZPOC-13 QTZPOC-13: Introduce joblib fallback for MLflow model loading Enhances model loading resilience by attempting to deserialize models directly with joblib if MLflow's native flavor loaders fail. This mechanism handles cases where .pkl files are present in the artifact. Adds joblib as a core dependency and includes new end-to-end and unit tests to validate fallback behavior, covering success, no-fallback, and double-failure scenarios. Also updates .gitignore for .claude/ files and incorporates minor refactorings and linting adjustments across E2E tests and utility functions. --- .gitignore | 3 +- e2e/conftest.py | 242 +++++++---- e2e/helpers.py | 19 +- e2e/opc_test_server.py | 7 +- e2e/test_child_workflows_e2e.py | 4 +- e2e/test_minio_offload.py | 24 +- e2e/test_opc_real_server.py | 29 +- e2e/test_predictions_batch_format_export.py | 98 +++-- e2e/test_predictions_batch_joblib_fallback.py | 186 +++++++++ e2e/test_predictions_batch_main_workflow.py | 3 +- ...st_predictions_batch_prediction_process.py | 14 +- .../utils/repository/model_repository.py | 203 ++++++++- pyproject.toml | 5 + requirements-light.txt | 1 + requirements-local.txt | 4 +- requirements.txt | 1 + .../utils/repository/test_model_repository.py | 395 +++++++++++++++++- 17 files changed, 1063 insertions(+), 175 deletions(-) create mode 100644 e2e/test_predictions_batch_joblib_fallback.py diff --git a/.gitignore b/.gitignore index c849e0a..3bd4a6b 100644 --- a/.gitignore +++ b/.gitignore @@ -55,4 +55,5 @@ mlruns/ relatorio* openspec/ -.cursor/ \ No newline at end of file +.cursor/ +.claude/ \ No newline at end of file diff --git a/e2e/conftest.py b/e2e/conftest.py index cc49d9b..c9c4504 100644 --- a/e2e/conftest.py +++ b/e2e/conftest.py @@ -12,28 +12,35 @@ sys.modules.setdefault('sientia', MagicMock()) sys.modules.setdefault('sientia.ModelAnalysis', _model_analysis_module) from io import BytesIO +import joblib import pandas as pd import pytest import pytest_asyncio +from sientia_do.notifications.handlers import CoreNotificationHandler +from sientia_do.observability.metrics_controller import MetricsController +from sklearn.compose import ColumnTransformer +from sklearn.linear_model import LinearRegression +from sklearn.pipeline import Pipeline from sqlalchemy import create_engine, text -from testcontainers.minio import MinioContainer -from testcontainers.postgres import PostgresContainer from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker +from testcontainers.minio import MinioContainer +from testcontainers.postgres import PostgresContainer from e2e.opc_test_server import OpcE2ETestServer from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch -from laborious.workflows.sub_workflows.prediction_process import PredictionProcess from laborious.workflows.sub_workflows.format_and_export_prediction import FormatAndExportPrediction -from sientia_do.notifications.handlers import CoreNotificationHandler -from sientia_do.observability.logger import Logger -from sientia_do.observability.metrics_controller import MetricsController +from laborious.workflows.sub_workflows.prediction_process import PredictionProcess # Test constants TEST_MONGODB_CONNECTION_STRING = 'mongodb://localhost:27017' TEST_DATABASE_NAME = 'test_db' +# Dedicated model whose MLflow flavor loaders fail, so the joblib fallback is exercised. +JOBLIB_FALLBACK_MODEL_NAME = 'joblib_fallback_model' +JOBLIB_FALLBACK_FLAVOR_ERROR = 'flavor loader cannot deserialize this artifact' + @pytest_asyncio.fixture(scope='session') def minio_container(): @@ -50,7 +57,7 @@ def minio_container(): def postgres_container(): """ Create a PostgreSQL container using testcontainers. - + This fixture creates a real PostgreSQL database in a Docker container that will be used for all tests in the session. """ @@ -64,21 +71,21 @@ def postgres_container(): def postgres_engine(postgres_container): """ Create SQLAlchemy engine for PostgreSQL test database. - + This fixture creates a connection to the PostgreSQL container created by the postgres_container fixture. """ engine = create_engine(postgres_container.get_connection_url()) - + yield engine - + engine.dispose() def _create_schema_and_tables(engine): """ Helper function to create schema and tables in the given engine. - + Creates predictions_schema with: - laborious_data: Input data table for queries - predictions: Output predictions table @@ -87,8 +94,8 @@ def _create_schema_and_tables(engine): # Use begin() to ensure transaction is properly committed with engine.begin() as conn: # Create predictions_schema - conn.execute(text("CREATE SCHEMA IF NOT EXISTS predictions_schema")) - + conn.execute(text('CREATE SCHEMA IF NOT EXISTS predictions_schema')) + # Create laborious_data table (input data from sensors) create_laborious_data_sql = """ CREATE TABLE IF NOT EXISTS predictions_schema.laborious_data ( @@ -102,7 +109,7 @@ def _create_schema_and_tables(engine): ); """ conn.execute(text(create_laborious_data_sql)) - + # Create predictions table create_predictions_sql = """ CREATE TABLE if not exists predictions_schema.predictions ( @@ -119,7 +126,7 @@ def _create_schema_and_tables(engine): ); """ conn.execute(text(create_predictions_sql)) - + # Create transformed_data table create_transformed_sql = """ CREATE TABLE IF NOT EXISTS predictions_schema.transformed_data ( @@ -139,7 +146,7 @@ def _create_schema_and_tables(engine): def setup_postgres_schema_and_tables(postgres_engine): """ Automatically create necessary schema and tables before each test. - + This fixture runs automatically (autouse=True) and ensures that the predictions_schema and tables exist with the correct structure. """ @@ -150,35 +157,22 @@ def setup_postgres_schema_and_tables(postgres_engine): @pytest_asyncio.fixture def mock_logger(): """Mock logger for testing.""" + def message(message): - print(f"[LOG] {message}") - def custom_message(message, _metadata={}): - print(f"[LOG] {message}") + print(f'[LOG] {message}') + + def custom_message(message, _metadata=None): + print(f'[LOG] {message}') + logger = MagicMock() - logger.info = MagicMock( - side_effect=message - ) - logger.debug = MagicMock( - side_effect=message - ) - logger.error = MagicMock( - side_effect=message - ) - logger.warning = MagicMock( - side_effect=message - ) - logger.custom_info = MagicMock( - side_effect=custom_message - ) - logger.custom_debug = MagicMock( - side_effect=custom_message - ) - logger.custom_error = MagicMock( - side_effect=custom_message - ) - logger.custom_warning = MagicMock( - side_effect=custom_message - ) + logger.info = MagicMock(side_effect=message) + logger.debug = MagicMock(side_effect=message) + logger.error = MagicMock(side_effect=message) + logger.warning = MagicMock(side_effect=message) + logger.custom_info = MagicMock(side_effect=custom_message) + logger.custom_debug = MagicMock(side_effect=custom_message) + logger.custom_error = MagicMock(side_effect=custom_message) + logger.custom_warning = MagicMock(side_effect=custom_message) return logger @@ -186,24 +180,24 @@ def mock_logger(): def mock_mongo_client(): """ Mock MongoDB client to avoid real connections. - + This fixture mocks the pymongo.MongoClient used by CoreNotificationHandler, allowing us to use a real NotificationHandler instance without connecting to MongoDB. """ mock_client = MagicMock() mock_db = MagicMock() mock_collection = MagicMock() - + # Configure the mock chain: client[database] -> db[collection] -> collection mock_client.__getitem__.return_value = mock_db mock_db.__getitem__.return_value = mock_collection - + # Mock server_info() to avoid connection attempts mock_client.server_info = MagicMock() - + # Mock insert_one for notifications mock_collection.insert_one = MagicMock() - + return mock_client @@ -225,7 +219,7 @@ def notification_inserts(mock_mongo_client): def notification_handler(mock_logger, mock_mongo_client): """ Create a real NotificationHandler instance with mocked MongoDB client. - + This fixture creates a real CoreNotificationHandler instance but mocks the underlying MongoDB connection to avoid real database connections. """ @@ -258,7 +252,7 @@ def mock_minio_repository(): parquet_buffer = BytesIO() parquet_df.to_parquet(parquet_buffer, engine='pyarrow', index=True) parquet_bytes = parquet_buffer.getvalue() - + # sientia_do MinioRepository API mock_repo.bucket = 'test-bucket' mock_repo.upload_file = AsyncMock( @@ -273,7 +267,7 @@ def mock_minio_repository(): mock_repo.list_objects = AsyncMock(return_value=[]) mock_repo.delete_file = AsyncMock() mock_repo.close = MagicMock() - + return mock_repo @@ -295,13 +289,12 @@ def mock_pi_web_api_repository(): mock_repo.close = MagicMock() return mock_repo + @pytest_asyncio.fixture def mock_opc_repository(): """Mock OPC repository for OPC operations.""" mock_repo = MagicMock() - mock_repo.write_data = AsyncMock( - return_value=(True, {'response_time': 0.1}) - ) + mock_repo.write_data = AsyncMock(return_value=(True, {'response_time': 0.1})) mock_repo.disconnect = AsyncMock() return mock_repo @@ -318,10 +311,13 @@ async def opc_e2e_server(): finally: await server.stop() + @pytest_asyncio.fixture def patch_create_engine(postgres_engine): """Patch create_engine to return test postgres_engine.""" - with patch('sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine): + with patch( + 'sientia_do.temporal.activities.postgres.create_engine', return_value=postgres_engine + ): yield @@ -329,44 +325,54 @@ def patch_create_engine(postgres_engine): def patch_minio_repository(mock_minio_repository): """Patch MinioRepository to return mock.""" # Patch where Activities resolves the symbol (import binds the original class). - with patch('laborious.activities.activities.MinioRepository', return_value=mock_minio_repository): + with patch( + 'laborious.activities.activities.MinioRepository', return_value=mock_minio_repository + ): yield + @pytest_asyncio.fixture def patch_pi_web_api_repository(mock_pi_web_api_repository): """Patch MLflowRepository to return mock.""" with patch('laborious.activities.api.PIWebAPIClient', return_value=mock_pi_web_api_repository): yield + @pytest_asyncio.fixture def mock_mlflow_models(): """Create mock models for MLflow load_model methods.""" # Mock transform model - returns DataFrame with same index as input mock_transform_model = MagicMock() + def mock_transform_predict(data): num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1 print(data.to_csv()) print(data.index) - result = pd.DataFrame({ - 'feature_1': [0.234] * num_rows, - 'feature_2': [0.783] * num_rows, - }) + result = pd.DataFrame( + { + 'feature_1': [0.234] * num_rows, + 'feature_2': [0.783] * num_rows, + } + ) result.index = data.index return result + mock_transform_model.predict = MagicMock(side_effect=mock_transform_predict) - + # Mock predict model - returns array/list of predictions mock_predict_model = MagicMock() + def mock_predict_predict(data): num_rows = max(len(data), 1) if hasattr(data, '__len__') else 1 return [0.5] * num_rows + mock_predict_model.predict = MagicMock(side_effect=mock_predict_predict) - + # Mock PyFuncModel for compressed models mock_pyfunc_model = MagicMock() mock_pyfunc_model._model_impl = MagicMock() mock_pyfunc_model._model_impl.python_model = mock_transform_model - + return { 'transform_model': mock_transform_model, 'predict_model': mock_predict_model, @@ -374,50 +380,130 @@ def mock_mlflow_models(): } +def _build_joblib_fallback_estimator(): + """ + Build a real sklearn estimator for the joblib fallback artifact. + + `request_predict` appends a string `timestamp` column to the transformed + data, so the estimator selects the transform output columns explicitly and + drops the rest. + """ + features = ['feature_1', 'feature_2'] + estimator = Pipeline( + steps=[ + ( + 'select', + ColumnTransformer([('features', 'passthrough', features)], remainder='drop'), + ), + ('regressor', LinearRegression()), + ] + ) + train = pd.DataFrame( + { + 'feature_1': [0.0, 1.0, 2.0, 3.0], + 'feature_2': [0.0, 1.0, 0.0, 1.0], + 'timestamp': ['2024-01-01 12:00:00'] * 4, + } + ) + estimator.fit(train, train['feature_1'] * 2 + train['feature_2'] * 3) + return estimator + + +@pytest.fixture +def joblib_fallback_artifact(tmp_path): + """ + Real artifact directory holding a joblib-compressed sklearn estimator. + + Mirrors what the joblib fallback finds after downloading `prediction_model`: + a top-level `model.pkl` written by `joblib.dump(..., compress=3)` plus an + `MLmodel` placeholder. + + Return: + dict with 'path' (the artifact directory) and 'estimator' (the object + that was serialized, so tests can compute the expected prediction). + """ + artifact_dir = tmp_path / 'prediction_model' + artifact_dir.mkdir() + + estimator = _build_joblib_fallback_estimator() + joblib.dump(estimator, artifact_dir / 'model.pkl', compress=3) + (artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n') + + return {'path': str(artifact_dir), 'estimator': estimator} + + +@pytest.fixture +def joblib_fallback_artifact_without_pickle(tmp_path): + """Artifact directory with no top-level `.pkl`, so the joblib fallback also fails.""" + artifact_dir = tmp_path / 'prediction_model_without_pickle' + artifact_dir.mkdir() + (artifact_dir / 'MLmodel').write_text('flavors:\n sklearn: {}\n') + return str(artifact_dir) + + @pytest_asyncio.fixture -def patch_mlflow(mock_mlflow_models): +def patch_mlflow(mock_mlflow_models, joblib_fallback_artifact): """Patch mlflow module in repository with load_model mocks.""" mock_mlflow = MagicMock() - + # Mock sklearn.load_model def mock_sklearn_load_model(model_uri): + if JOBLIB_FALLBACK_MODEL_NAME in model_uri: + raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR) if 'data_model' in model_uri or 'transform' in model_uri.lower(): return mock_mlflow_models['transform_model'] return mock_mlflow_models['predict_model'] + mock_mlflow.sklearn = MagicMock() mock_mlflow.sklearn.load_model = MagicMock(side_effect=mock_sklearn_load_model) - + # Mock pyfunc.load_model def mock_pyfunc_load_model(model_uri): + if JOBLIB_FALLBACK_MODEL_NAME in model_uri: + raise RuntimeError(JOBLIB_FALLBACK_FLAVOR_ERROR) if 'artifacts' in model_uri or 'tmp' in model_uri: return mock_mlflow_models['pyfunc_model'] if 'data_model' in model_uri or 'transform' in model_uri.lower(): return mock_mlflow_models['transform_model'] return mock_mlflow_models['predict_model'] + mock_mlflow.pyfunc = MagicMock() mock_mlflow.pyfunc.load_model = MagicMock(side_effect=mock_pyfunc_load_model) - + # Mock pytorch.load_model mock_mlflow.pytorch = MagicMock() mock_mlflow.pytorch.load_model = MagicMock(return_value=mock_mlflow_models['predict_model']) - + # Mock other mlflow methods that might be called mock_mlflow.set_tracking_uri = MagicMock() - mock_mlflow.get_run = MagicMock(return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id'))) + mock_mlflow.get_run = MagicMock( + return_value=MagicMock(info=MagicMock(artifact_uri='mlflow-artifacts:/test_run_id')) + ) mock_mlflow.tracking = MagicMock() - mock_mlflow.tracking.MlflowClient = MagicMock(return_value=MagicMock( + mock_mlflow_client = MagicMock( search_registered_models=MagicMock(return_value=[MagicMock(name='test_model')]), - search_model_versions=MagicMock(return_value=[MagicMock( - current_stage='Production', - version='1', - source='runs:/artifacts/test_run_id' - )]) - )) - + search_model_versions=MagicMock( + return_value=[ + MagicMock( + current_stage='Production', version='1', source='runs:/artifacts/test_run_id' + ) + ] + ), + ) + # Only the joblib fallback calls download_artifacts directly on the client. + mock_mlflow_client.download_artifacts = MagicMock(return_value=joblib_fallback_artifact['path']) + mock_mlflow.tracking.MlflowClient = MagicMock(return_value=mock_mlflow_client) + with patch('laborious.utils.repository.model_repository.mlflow', new=mock_mlflow): yield mock_mlflow +@pytest.fixture +def mlflow_client(patch_mlflow): + """The MlflowClient mock the repository holds as `self.client`.""" + return patch_mlflow.tracking.MlflowClient.return_value + + @pytest_asyncio.fixture(scope='function') async def test_activities( postgres_engine, @@ -430,7 +516,7 @@ async def test_activities( patch_minio_repository, patch_mlflow, patch_pi_web_api_repository, - mock_opc_repository + mock_opc_repository, ): """ Create Activities instance with test dependencies. diff --git a/e2e/helpers.py b/e2e/helpers.py index bb404ca..2919ce9 100644 --- a/e2e/helpers.py +++ b/e2e/helpers.py @@ -11,7 +11,9 @@ from sqlalchemy import text from sqlalchemy.engine import Engine -async def start_and_await_workflow(client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0): +async def start_and_await_workflow( + client, workflow_run, input_data: dict, workflow_id: str, timeout: float = 60.0 +): """ Start a workflow and wait for its result. @@ -44,7 +46,9 @@ def insert_sample_data(postgres_engine: Engine, model_id: int, values: list[Any] values: Per-sensor values; use string 'NULL' for SQL NULL. """ with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}') + ) values_sql = [] for i, value in enumerate(values): values_sql.append(f""" @@ -79,7 +83,6 @@ def assert_prediction( comments: Expected exact comments string (optional). comments_contains: Substring expected in comments when queued (optional). """ - import pytest with postgres_engine.connect() as conn: result_query = conn.execute( @@ -90,7 +93,9 @@ def assert_prediction( ) ) prediction_rows = result_query.fetchall() - assert len(prediction_rows) == 1, f'Expected one prediction record, got {len(prediction_rows)}' + assert len(prediction_rows) == 1, ( + f'Expected one prediction record, got {len(prediction_rows)}' + ) row = prediction_rows[0] assert row[0] == model_id, f'Expected model_id={model_id}, got {row[0]}' assert row[1] == prediction or Decimal(str(row[1])) == Decimal(str(prediction)), ( @@ -99,7 +104,9 @@ def assert_prediction( assert row[2] == prediction_confidence or Decimal(str(row[2])) == Decimal( str(prediction_confidence) ), f'Expected prediction_confidence={prediction_confidence}, got {row[2]}' - assert row[3] == prediction_status, f"Expected prediction_status='{prediction_status}', got {row[3]}" + assert row[3] == prediction_status, ( + f"Expected prediction_status='{prediction_status}', got {row[3]}" + ) if comments is not None: assert row[4] == comments, f"Expected comments='{comments}', got {row[4]}" if comments_contains is not None: @@ -135,7 +142,6 @@ def assert_continue( def assert_stop(postgres_engine: Engine, model_id: int) -> None: """Assert no prediction rows for model_id.""" - import pytest with postgres_engine.connect() as conn: result_query = conn.execute( @@ -156,7 +162,6 @@ def assert_repeat(postgres_engine: Engine, model_id: int, last_prediction: tuple model_id: Model id. last_prediction: Tuple (model_id, prediction, confidence, status) to match both rows. """ - import pytest with postgres_engine.connect() as conn: result_query = conn.execute( diff --git a/e2e/opc_test_server.py b/e2e/opc_test_server.py index 69b02c8..8f94eea 100644 --- a/e2e/opc_test_server.py +++ b/e2e/opc_test_server.py @@ -147,8 +147,11 @@ class OpcE2ETestServer: value = await self._confidence_node.read_value() return float(value) - async def _pre_write_callback(self, _event, _service) -> None: - if self._session_bad_on_write: + async def _pre_write_callback(self, event, _service) -> None: + # Only client writes are faulted. asyncua's own clock task writes CurrentTime + # through this same PreWrite dispatch, and raising there kills its time_task; + # the stored exception then resurfaces as a teardown failure in server.stop(). + if self._session_bad_on_write and getattr(event, 'is_external', False): raise ServiceError(ua.StatusCodes.BadSessionIdInvalid) diff --git a/e2e/test_child_workflows_e2e.py b/e2e/test_child_workflows_e2e.py index 52a4b3f..34d97fd 100644 --- a/e2e/test_child_workflows_e2e.py +++ b/e2e/test_child_workflows_e2e.py @@ -26,7 +26,9 @@ async def test_format_and_export_prediction_default_path_e2e( client = temporal_test_env.client model_id = 401 with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) metadata = { 'metadata': { diff --git a/e2e/test_minio_offload.py b/e2e/test_minio_offload.py index f8c0122..dd303f1 100644 --- a/e2e/test_minio_offload.py +++ b/e2e/test_minio_offload.py @@ -29,7 +29,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket( """ model_id = 501 with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}') + ) insert_sample_data(postgres_engine, model_id, [1.0, 2.0]) metadata = { @@ -53,7 +55,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket( } ) assert payload.object_key, 'offloaded payload must reference a MinIO object' - assert payload.data is None or payload.data == {}, 'large payloads should not inline tabular dict' + assert payload.data is None or payload.data == {}, ( + 'large payloads should not inline tabular dict' + ) df = await payload.retrieve(test_activities_real_minio.minio_repository, metadata['metadata']) assert len(df) >= 1 @@ -61,7 +65,9 @@ async def test_load_query_with_minio_offload_writes_object_to_bucket( client = minio_container.get_client() listed = list(client.list_objects('test-bucket', recursive=True)) names = [getattr(o, 'object_name', None) or getattr(o, '_object_name', '') for o in listed] - assert any(n and 'prediction_datasets' in n for n in names), f'unexpected object listing: {names!r}' + assert any(n and 'prediction_datasets' in n for n in names), ( + f'unexpected object listing: {names!r}' + ) @pytest.mark.asyncio @@ -77,9 +83,15 @@ async def test_predictions_batch_with_minio_offload_path( """ model_id = 502 with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) - conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) - conn.execute(text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}') + ) + conn.execute( + text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) + conn.execute( + text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}') + ) insert_sample_data(postgres_engine, model_id, [10.0, 20.0, 30.0]) input_data = { diff --git a/e2e/test_opc_real_server.py b/e2e/test_opc_real_server.py index 200a8a4..7445656 100644 --- a/e2e/test_opc_real_server.py +++ b/e2e/test_opc_real_server.py @@ -11,7 +11,12 @@ import pytest from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker -from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow +from e2e.helpers import ( + assert_prediction, + insert_sample_data, + make_workflow_id, + start_and_await_workflow, +) from e2e.opc_test_server import UNKNOWN_NODE_ID, OpcE2ETestServer, build_opc_output_config from e2e.test_predictions_batch_format_export import get_base_input_data from laborious.activities.activities import Activities @@ -20,16 +25,24 @@ from laborious.utils.repository.opc_repository import OpcRepository from laborious.workflows.predictions_batch import PredictionsBatch -async def _slow_reconnect_under_lock(repo: OpcRepository, hold_seconds: float = 0.75) -> None: +async def _reconnect_under_held_lock( + repo: OpcRepository, acquired: asyncio.Event, release: asyncio.Event +) -> None: """ - Hold the connection lock briefly so concurrent writes see reconnect_in_progress. + Hold the connection lock until released, then reconnect. + + The lock must stay held for the whole workflow, and how long the workflow takes to + reach the OPC write is not bounded — a fixed sleep here lets `_reconnect_locked` + finish first, which sets `_session_ready` again and the write then succeeds. Args: repo (OpcRepository): Connected repository. - hold_seconds (float): Time to keep the lock before reconnecting. + acquired (asyncio.Event): Set once the lock is held, before any write can race. + release (asyncio.Event): Awaited while holding the lock; set by the caller. """ async with repo._connection_lock: - await asyncio.sleep(hold_seconds) + acquired.set() + await release.wait() await repo._reconnect_locked() @@ -170,7 +183,10 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server( repo = test_activities_real_opc.opc_repository['1'] repo._session_ready.clear() - reconnect_task = asyncio.create_task(_slow_reconnect_under_lock(repo)) + acquired = asyncio.Event() + release = asyncio.Event() + reconnect_task = asyncio.create_task(_reconnect_under_held_lock(repo, acquired, release)) + await acquired.wait() input_data = get_base_input_data(model_id) input_data['opc_output_config'] = build_opc_output_config(opc_e2e_server.node_ids) @@ -184,6 +200,7 @@ async def test_scenario_3_2_5_opc_write_blocked_during_reconnect_real_server( make_workflow_id('test-opc-real-reconnect-block'), ) finally: + release.set() await reconnect_task assert_prediction( diff --git a/e2e/test_predictions_batch_format_export.py b/e2e/test_predictions_batch_format_export.py index 5582d0c..cc338c4 100644 --- a/e2e/test_predictions_batch_format_export.py +++ b/e2e/test_predictions_batch_format_export.py @@ -4,7 +4,7 @@ End-to-end tests for PredictionsBatch workflow - Format and Export scenarios. from decimal import Decimal from typing import Any, cast -from unittest.mock import ANY, AsyncMock, call +from unittest.mock import AsyncMock, call import pytest from sientia_do.notifications.models import NotificationLevel @@ -12,7 +12,12 @@ from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker -from e2e.helpers import assert_prediction, insert_sample_data, make_workflow_id, start_and_await_workflow +from e2e.helpers import ( + assert_prediction, + insert_sample_data, + make_workflow_id, + start_and_await_workflow, +) from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch @@ -46,7 +51,8 @@ base_input_data = { 'datetime_columns': ['timestamp', 'created_at'], } -base_query = "SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}" +base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + def get_base_input_data(model_id): return { @@ -79,8 +85,12 @@ async def test_scenario_3_1_1_default_prediction_export( model_id = 311 with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}")) - conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) + conn.execute( + text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) + conn.execute( + text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}') + ) insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) input_data = get_base_input_data(model_id) @@ -176,7 +186,9 @@ async def test_scenario_3_1_1_default_prediction_export( with postgres_engine.connect() as conn: tf_count = conn.execute( - text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}") + text( + f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}' + ) ).scalar() assert tf_count == 0, 'transform export must be skipped when path_flag is set' @@ -190,9 +202,6 @@ async def test_scenario_3_1_1_default_prediction_export( ) - - - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_1_2_export_with_opc_only( @@ -203,17 +212,17 @@ async def test_scenario_3_1_2_export_with_opc_only( ): """ Scenario 3.1.2: Export with OPC only - + Description: Export to PostgreSQL and OPC server only (no PI Web API). - + Expected Behavior: - Normal formatting - PostgreSQL export executed - OPC export executed - PI Web API activity skipped - Metrics written with OPC metrics - + Assertions: - PI Web API activity NOT called - OPC activity called @@ -280,7 +289,6 @@ async def test_scenario_3_1_2_export_with_opc_only( assert_prediction(postgres_engine, model_id) - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_1_3_export_with_pi_web_api_only( @@ -291,17 +299,17 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( ): """ Scenario 3.1.3: Export with PI Web API only - + Description: Export to PostgreSQL and PI Web API only (no OPC). - + Expected Behavior: - Normal formatting - PostgreSQL export executed - PI Web API export executed - OPC activity skipped - Metrics written without OPC metrics - + Assertions: - OPC activity NOT called - PI Web API activity called @@ -364,7 +372,6 @@ async def test_scenario_3_1_3_export_with_pi_web_api_only( assert_prediction(postgres_engine, model_id) - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_1_4_export_without_optional_outputs( @@ -375,16 +382,16 @@ async def test_scenario_3_1_4_export_without_optional_outputs( ): """ Scenario 3.1.4: Export Without Optional Outputs - + Description: Export only to PostgreSQL (no OPC or PI Web API). - + Expected Behavior: - Normal formatting - Only PostgreSQL export executed - OPC and PI Web API activities skipped - Metrics written without OPC metrics - + Assertions: - PI Web API activity NOT called - OPC activity NOT called @@ -412,7 +419,6 @@ async def test_scenario_3_1_4_export_without_optional_outputs( assert_prediction(postgres_engine, model_id) - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_1_5_export_without_transformed_data( @@ -423,15 +429,15 @@ async def test_scenario_3_1_5_export_without_transformed_data( ): """ Scenario 3.1.5: Export Without Transformed Data - + Description: Only prediction exported, no transform table. - + Expected Behavior: - Only prediction formatted and exported - Transform export skipped - Single PostgreSQL write - + Assertions: - format_transformed_data NOT called - One PostgreSQL export @@ -443,7 +449,9 @@ async def test_scenario_3_1_5_export_without_transformed_data( insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) with postgres_engine.begin() as conn: - conn.execute(text(f"DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}")) + conn.execute( + text(f'DELETE FROM predictions_schema.transformed_data WHERE model_id = {model_id}') + ) input_data = get_base_input_data(model_id) input_data['save_transform'] = False # Don't save transformed data @@ -533,15 +541,16 @@ async def test_scenario_3_1_5_export_without_transformed_data( with postgres_engine.connect() as conn: result_query = conn.execute( - text(f"SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}") + text( + f'SELECT COUNT(*) FROM predictions_schema.transformed_data WHERE model_id = {model_id}' + ) ) count = result_query.scalar() - assert count == 0, f"Expected transform table to be empty, but found {count} records" + assert count == 0, f'Expected transform table to be empty, but found {count} records' assert_prediction(postgres_engine, model_id) - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_2_1_pi_web_api_write_error( @@ -565,7 +574,8 @@ async def test_scenario_3_2_1_pi_web_api_write_error( insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) test_activities.pi_web_api_client.write_value.side_effect = Exception( - "PI Web API service unavailable") + 'PI Web API service unavailable' + ) input_data = get_base_input_data(model_id) input_data['pi_web_api_output_config'] = { @@ -593,14 +603,14 @@ async def test_scenario_3_2_1_pi_web_api_write_error( ) assert_prediction( - postgres_engine, model_id, + postgres_engine, + model_id, prediction_confidence=13, comments='PI Web API service unavailable', ) assert notification_inserts.call_count >= 1 - @pytest.mark.asyncio @pytest.mark.integration async def test_scenario_3_2_2_opc_write_error( @@ -622,14 +632,17 @@ async def test_scenario_3_2_2_opc_write_error( insert_sample_data(postgres_engine, model_id, [23.5, 78.2]) opc_write_data = cast(Any, test_activities.opc_repository['1'].write_data) - opc_write_data.return_value = (False, { - 'notification_id': 'OPC_WRITE_DATA_ERROR_1', - 'message': 'OPC server unavailable', - 'block': 'opc_repository', - 'level': NotificationLevel.ERROR, - 'attachment_content': 'OPC server unavailable', - }) - + opc_write_data.return_value = ( + False, + { + 'notification_id': 'OPC_WRITE_DATA_ERROR_1', + 'message': 'OPC server unavailable', + 'block': 'opc_repository', + 'level': NotificationLevel.ERROR, + 'attachment_content': 'OPC server unavailable', + }, + ) + input_data = get_base_input_data(model_id) input_data['opc_output_config'] = { '1': { @@ -656,7 +669,8 @@ async def test_scenario_3_2_2_opc_write_error( ) assert_prediction( - postgres_engine, model_id, + postgres_engine, + model_id, prediction_confidence=12, comments='Some data could not be written to OPC servers', ) @@ -779,8 +793,8 @@ async def test_scenario_3_2_3_pi_web_api_partial_write_error( ) assert_prediction( - postgres_engine, model_id, + postgres_engine, + model_id, prediction_confidence=13, comments="The number of written tags does not match the number of tag names: Expected ['tag_1', 'tag_3'] tags, but ['tag_1'] tags were written.", ) - diff --git a/e2e/test_predictions_batch_joblib_fallback.py b/e2e/test_predictions_batch_joblib_fallback.py new file mode 100644 index 0000000..a807c3e --- /dev/null +++ b/e2e/test_predictions_batch_joblib_fallback.py @@ -0,0 +1,186 @@ +""" +End-to-end tests for the joblib fallback on the MLflow model load paths. + +`joblib.load` is NOT mocked here: the artifact fixture holds a real sklearn +estimator serialized with `joblib.dump(..., compress=3)`, so a successful run +proves the fallback really deserializes a compressed pickle and serves it. +""" + +import pandas as pd +import pytest +from sqlalchemy import text +from temporalio.testing import WorkflowEnvironment +from temporalio.worker import Worker + +from e2e.conftest import JOBLIB_FALLBACK_FLAVOR_ERROR, JOBLIB_FALLBACK_MODEL_NAME +from e2e.helpers import ( + assert_prediction, + assert_stop, + insert_sample_data, + make_workflow_id, + start_and_await_workflow, +) +from laborious.activities.activities import Activities +from laborious.workflows.predictions_batch import PredictionsBatch + +base_query = 'SELECT timestamp, variable, value, created_at FROM predictions_schema.laborious_data WHERE model_id = {model_id}' + + +def build_input_data(model_id: int, model_name: str) -> dict: + """Workflow input for a plain transform + predict run of a given model.""" + return { + 'schedule_name': 'test-schedule', + 'model_name': model_name, + 'model_id': model_id, + 'query': base_query.format(model_id=model_id), + 'schema': 'predictions_schema', + 'table_name': 'predictions', + 'transform_table_name': 'transformed_data', + 'input_filters': { + 'SPECIFIC_VARIABLES_NULL_VALUES': { + 'POLICY': 'CONTINUE', + 'CONFIG': {'variables': ['sensor_1']}, + }, + }, + 'mlflow_transform_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'mlflow_predict_filters': { + 'API_ERROR': {'POLICY': 'STOP', 'CONFIG': {}}, + }, + 'path_priority': ['STOP', 'CONTINUE', 'REPEAT'], + 'opc_output_config': {}, + 'pi_web_api_output_config': {}, + 'save_transform': True, + 'prediction_store_policy': 'lts:1', + 'model_config': { + 'retention_minutes': 0, + 'transform_flavor': 'sklearn', + 'predict_flavor': 'sklearn', + }, + 'datetime_columns': ['timestamp', 'created_at'], + } + + +def expected_fallback_prediction(estimator) -> float: + """Prediction the fixture estimator produces for the mocked transform output.""" + features = pd.DataFrame( + { + 'feature_1': [0.234], + 'feature_2': [0.783], + 'timestamp': ['2024-01-01 12:00:00'], + } + ) + return float(estimator.predict(features)[0]) + + +def assert_single_prediction_close(postgres_engine, model_id, expected, tolerance=1e-9): + """Assert exactly one Good prediction row whose value matches `expected`.""" + with postgres_engine.connect() as conn: + rows = conn.execute( + text( + 'SELECT prediction, prediction_status ' + 'FROM predictions_schema.predictions ' + f'WHERE model_id = {model_id}' + ) + ).fetchall() + + assert len(rows) == 1, f'Expected one prediction record, got {len(rows)}' + assert abs(float(rows[0][0]) - expected) < tolerance, ( + f'Expected prediction close to {expected}, got {rows[0][0]}' + ) + assert rows[0][1] == 'Good', f"Expected prediction_status='Good', got {rows[0][1]}" + + +def logged_messages(mock_logger, level: str) -> str: + """Concatenate every message logged at a given custom level.""" + calls = getattr(mock_logger, f'custom_{level}').call_args_list + return ' '.join(str(logged.args[0]) for logged in calls) + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_prediction_workflow_recovered_by_joblib_fallback( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + joblib_fallback_artifact, + mlflow_client, + mock_logger, +): + """Flavor loader fails, joblib fallback loads the real compressed pickle, prediction lands.""" + client = temporal_test_env.client + model_id = 271 + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME) + + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback') + ) + + assert_single_prediction_close( + postgres_engine, + model_id, + expected_fallback_prediction(joblib_fallback_artifact['estimator']), + ) + + assert mlflow_client.download_artifacts.call_count == 1 + assert mlflow_client.download_artifacts.call_args.args[1] == 'prediction_model' + + assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged_messages(mock_logger, 'warning') + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_prediction_workflow_without_fallback_is_unchanged( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + mlflow_client, + mock_mlflow_models, +): + """A model whose flavor loader works keeps the current result and downloads nothing.""" + client = temporal_test_env.client + model_id = 272 + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = build_input_data(model_id, 'test_model') + + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-no-fallback') + ) + + assert_prediction(postgres_engine, model_id, prediction=0.5) + mock_mlflow_models['predict_model'].predict.assert_called() + mlflow_client.download_artifacts.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.integration +async def test_prediction_workflow_reports_both_causes_when_fallback_fails( + temporal_test_env: WorkflowEnvironment, + temporal_worker: Worker, + test_activities: Activities, + postgres_engine, + joblib_fallback_artifact_without_pickle, + mlflow_client, + mock_logger, +): + """Both attempts fail: workflow stops and the logged traceback carries both causes.""" + mlflow_client.download_artifacts.return_value = joblib_fallback_artifact_without_pickle + + client = temporal_test_env.client + model_id = 273 + insert_sample_data(postgres_engine, model_id, [60.0, 78.2]) + input_data = build_input_data(model_id, JOBLIB_FALLBACK_MODEL_NAME) + + await start_and_await_workflow( + client, PredictionsBatch.run, input_data, make_workflow_id('test-joblib-fallback-failed') + ) + + assert_stop(postgres_engine, model_id) + + logged = logged_messages(mock_logger, 'debug') + assert JOBLIB_FALLBACK_FLAVOR_ERROR in logged + assert 'No .pkl file found in prediction_model artifact' in logged diff --git a/e2e/test_predictions_batch_main_workflow.py b/e2e/test_predictions_batch_main_workflow.py index 3f8d37f..9719a9f 100644 --- a/e2e/test_predictions_batch_main_workflow.py +++ b/e2e/test_predictions_batch_main_workflow.py @@ -4,12 +4,11 @@ End-to-end tests for PredictionsBatch workflow - Main workflow scenarios. import asyncio +import pytest from sqlalchemy import text from temporalio.testing import WorkflowEnvironment from temporalio.worker import Worker -import pytest - from e2e.helpers import make_workflow_id, start_and_await_workflow from laborious.activities.activities import Activities from laborious.workflows.predictions_batch import PredictionsBatch diff --git a/e2e/test_predictions_batch_prediction_process.py b/e2e/test_predictions_batch_prediction_process.py index af097cd..6623c9e 100644 --- a/e2e/test_predictions_batch_prediction_process.py +++ b/e2e/test_predictions_batch_prediction_process.py @@ -3,7 +3,7 @@ End-to-end tests for PredictionsBatch workflow - Prediction Process scenarios. """ from decimal import Decimal -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock import numpy as np import pandas as pd @@ -69,7 +69,9 @@ def get_base_input_data(model_id): def insert_sample_prediction(postgres_engine, model_id): with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) insert_sql = f""" INSERT INTO predictions_schema.predictions (model_id, timestamp, prediction, prediction_confidence, prediction_status, comments, response_time) VALUES @@ -180,7 +182,9 @@ async def test_scenario_2_1_4_input_gate_repeat_without_prior_prediction( model_id = 214 insert_sample_data(postgres_engine, model_id, ['NULL', 78.2]) with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.predictions WHERE model_id = {model_id}') + ) input_data = get_base_input_data(model_id) input_data['input_filters']['SPECIFIC_VARIABLES_NULL_VALUES']['POLICY'] = 'REPEAT' await start_and_await_workflow( @@ -370,7 +374,9 @@ async def test_scenario_2_4_1_input_empty_data_stop( client = temporal_test_env.client model_id = 241 with postgres_engine.begin() as conn: - conn.execute(text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}')) + conn.execute( + text(f'DELETE FROM predictions_schema.laborious_data WHERE model_id = {model_id}') + ) input_data = get_base_input_data(model_id) input_data['input_filters'] = {'EMPTY_DATA': {'POLICY': 'STOP', 'CONFIG': {}}} await start_and_await_workflow( diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index 6a25b8c..d6dc2bc 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -21,10 +21,11 @@ import time import traceback from datetime import datetime, timedelta from io import StringIO -from os import environ, makedirs, path +from os import environ, listdir, makedirs, path from shutil import rmtree from typing import Any, Literal, overload +import joblib import mlflow import pandas as pd from mlflow.entities import Experiment @@ -351,6 +352,117 @@ class MLFlowRepository(SientiaMonitoring): self.info(f'Loaded dataframe from {model_name}:{artifact_path}', metadata) return dataframe + def _find_model_pickle( + self, local_dir: str, artifact_name: str, flavor_error: Exception + ) -> str: + """ + Locate the pickle file of a downloaded model artifact. + + Fallback path only: this is reached after the MLflow flavor loader has + already failed. The target artifact is always a single `.pkl` at the top + level of the artifact directory, so the listing does not recurse — + nested wrapper artifacts must never be deserialized in place of the + model. No other extension is considered and `MLmodel` is not parsed. + + Args: + local_dir (str): Local directory holding the downloaded artifact. + artifact_name (str): Artifact name, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader, + reported in the message so the original cause is not lost. + + Returns: + str: Path to the first top-level `.pkl` file, in name order. + + Raises: + ValueError: If the top level of `local_dir` has no `.pkl` file. + """ + entries = sorted(listdir(local_dir)) + pkl_files = [ + entry + for entry in entries + if entry.endswith('.pkl') and path.isfile(path.join(local_dir, entry)) + ] + + if not pkl_files: + raise ValueError( + f'No .pkl file found in {artifact_name} artifact. ' + f'Original error: {flavor_error}. ' + f'Top-level entries in {local_dir}: {entries}' + ) + + return path.join(local_dir, pkl_files[0]) + + def _load_model_pickle( + self, local_dir: str, artifact_name: str, flavor_error: Exception + ) -> Any: + """ + Deserialize a downloaded model artifact with `joblib`. + + Fallback path only. `joblib.load` reads the compressed `.pkl` produced + by `joblib.dump`, which is what the MLflow flavor loader fails to read. + The deserialized object is returned as is: no `._model_impl.python_model` + unwrap is applied, because the pickled object already is the model. + + Args: + local_dir (str): Local directory holding the downloaded artifact. + artifact_name (str): Artifact name, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader. + + Returns: + Any: The object deserialized from the artifact's pickle file. + + Raises: + ValueError: If the artifact directory has no top-level `.pkl` file. + Exception: Any error raised by `joblib.load` itself. + """ + pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error) + + self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}') + + return joblib.load(pkl_path) + + async def _download_and_load_pickle( + self, + model_name: str, + artifact_name: str, + flavor_error: Exception, + metadata: dict[str, Any], + ) -> Any: + """ + Download a model artifact and deserialize it with `joblib`. + + Fallback path only. The download runs without `dst_path`, so MLflow + writes to a per-call temporary directory: the shared `ARTIFACTS_PATH` + tree is never written to and `dowload_artifacts` — which removes its + destination before downloading — is deliberately not reused. The target + artifact is always a `.pkl` and the deserialized object is returned as + is. + + Args: + model_name (str): Registered model name. + artifact_name (str): Artifact to download, e.g. 'prediction_model'. + flavor_error (Exception): Error raised by the MLflow flavor loader. + metadata (dict[str, Any]): Metadata used for structured logging. + + Returns: + Any: The object deserialized from the artifact's pickle file. + + Raises: + ValueError: If the downloaded artifact has no top-level `.pkl` file. + Exception: Any error raised while resolving the Production run, + downloading the artifact, or deserializing the pickle. + """ + run_id = self.get_model_run_id(model_name=model_name, stage='Production') + + self.info( + f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}', + metadata, + ) + + local_dir = self.client.download_artifacts(run_id, artifact_name) + + return self._load_model_pickle(local_dir, artifact_name, flavor_error) + async def load_predict_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' ) -> Any: @@ -370,23 +482,43 @@ class MLFlowRepository(SientiaMonitoring): - The model is fetched from the "Production" stage of the MLflow Model Registry. - Warnings during the model loading process are suppressed. """ + if flavor not in ('pyfunc', 'sklearn', 'pytorch'): + raise ValueError(INVALID_FLAVOR_MESSAGE) + model_uri = f'models:/{model_name}/production' self.info(f'Loading prediction model {model_name} from {model_uri}') core_labels = self.get_core_labels(metadata, operation_type='load_predict_model') + fallback_labels = self.get_core_labels( + metadata, operation_type='load_predict_model_joblib_fallback' + ) start_time = time.time() try: if flavor == 'pyfunc': model = mlflow.pyfunc.load_model(model_uri) elif flavor == 'sklearn': model = mlflow.sklearn.load_model(model_uri) - elif flavor == 'pytorch': - model = mlflow.pytorch.load_model(model_uri) else: - raise ValueError(INVALID_FLAVOR_MESSAGE) - except Exception as e: - await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) - raise e + model = mlflow.pytorch.load_model(model_uri) + except Exception as flavor_error: + self.warning( + f'Flavor loader {flavor} failed for prediction model {model_name}, ' + f'trying joblib fallback. Original error: {flavor_error}', + metadata, + ) + try: + model = await self._download_and_load_pickle( + model_name, 'prediction_model', flavor_error, metadata + ) + except Exception as fallback_error: + await self.emit_metric( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels + ) + raise fallback_error from flavor_error + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + return model await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) @@ -415,25 +547,45 @@ class MLFlowRepository(SientiaMonitoring): model cannot be loaded. """ + if flavor not in ('sklearn', 'pyfunc', 'pytorch'): + raise ValueError(INVALID_FLAVOR_MESSAGE) + 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.info(f'Loading data model {model_name} from {model_uri}') core_labels = self.get_core_labels(metadata, operation_type='load_transform_model') + fallback_labels = self.get_core_labels( + metadata, operation_type='load_transform_model_joblib_fallback' + ) start_time = time.time() try: if flavor == 'sklearn': model = mlflow.sklearn.load_model(model_uri) elif flavor == 'pyfunc': model = mlflow.pyfunc.load_model(model_uri) - elif flavor == 'pytorch': - model = mlflow.pytorch.load_model(model_uri) else: - raise ValueError(INVALID_FLAVOR_MESSAGE) - except Exception as e: - await self.emit_metric(metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=core_labels) - raise e + model = mlflow.pytorch.load_model(model_uri) + except Exception as flavor_error: + self.warning( + f'Flavor loader {flavor} failed for data model {model_name}, ' + f'trying joblib fallback. Original error: {flavor_error}', + metadata, + ) + try: + model = await self._download_and_load_pickle( + model_name, 'data_model', flavor_error, metadata + ) + except Exception as fallback_error: + await self.emit_metric( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels + ) + raise fallback_error from flavor_error + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + return model await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) @@ -482,8 +634,29 @@ class MLFlowRepository(SientiaMonitoring): f'Model with type {model_type} and name {model_name} is compressed, loading from {artifact_path}' ) - raw_model = mlflow.pyfunc.load_model(artifact_path) - model = raw_model._model_impl.python_model + fallback_labels = self.get_core_labels( + metadata, operation_type='download_model_wrapper_joblib_fallback' + ) + start_time = time.time() + try: + raw_model = mlflow.pyfunc.load_model(artifact_path) + model = raw_model._model_impl.python_model + except Exception as flavor_error: + self.warning( + f'Wrapper loader failed for {model_type} model {model_name}, ' + f'trying joblib fallback. Original error: {flavor_error}', + metadata, + ) + try: + model = self._load_model_pickle(artifact_path, target, flavor_error) + except Exception as fallback_error: + await self.emit_metric( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels + ) + raise fallback_error from flavor_error + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) self.debug( f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata diff --git a/pyproject.toml b/pyproject.toml index db08645..4a90296 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,11 @@ ignore = [ "S105", # hardcoded passwords ok in tests "S106", # hardcoded passwords ok in tests ] +"e2e/conftest.py" = [ + # Imports below the sientia/ModelAnalysis stub are load-order critical: the stub + # must be in sys.modules before Activities is imported, so they cannot move up. + "E402", +] [tool.ruff.lint.mccabe] max-complexity = 15 diff --git a/requirements-light.txt b/requirements-light.txt index 302fb79..d5128de 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -16,3 +16,4 @@ shap pycurl scipy<1.14.0 scikit-learn==1.5.2 +joblib diff --git a/requirements-local.txt b/requirements-local.txt index 73a9a47..852328a 100644 --- a/requirements-local.txt +++ b/requirements-local.txt @@ -15,4 +15,6 @@ hyperopt shap pycurl scipy<1.14.0 -scikit-learn==1.5.2 \ No newline at end of file +scikit-learn==1.4.2 +pycaret==3.3.2 +joblib \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 4ad005d..8f994e9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,3 +16,4 @@ shap pycurl scipy<1.14.0 scikit-learn==1.5.2 +joblib diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 54a75bb..64c5e5f 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1,13 +1,20 @@ from datetime import UTC, datetime from unittest.mock import ANY, AsyncMock, MagicMock, call, patch +import joblib as joblib_lib import mlflow as mlflow_lib import numpy as np import pytest from pandas import DataFrame, Timestamp +from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ +from sklearn.linear_model import LinearRegression from laborious import metrics -from laborious.utils.repository.model_repository import MLFlowRepository, force_memory_release +from laborious.utils.repository.model_repository import ( + INVALID_FLAVOR_MESSAGE, + MLFlowRepository, + force_memory_release, +) @patch('laborious.utils.repository.model_repository.ctypes') @@ -608,32 +615,46 @@ def test_detect_and_parse_datetime_index_empty(mlflow_repository): assert response.equals(input_data) +# Each case is (index, fragments expected in the error message). +# +# The index is passed explicitly instead of through a dict-of-dict: building a +# DataFrame from nested dicts makes pandas sort the keys to infer the index, which a +# mixed-type index cannot do. +# +# Fragments, not whole messages: the parse branch appends pandas' own multi-line +# error, and the type branch appends the fully qualified class name, so neither is +# stable enough to compare literally. invalid_cases = [ ( - {'value': {'2024-01-01 12:00:00': 1, 2024: 2}}, - "Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Elements are , .", + ['2024-01-01 12:00:00', 2024], + ["Elements are ,"], ), ( - {'value': {'2024-01-01': 1, '2024-01-02': 2}}, - 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Unable to parse given date format', + ['2024-01-01', '2024-01-02'], + ['Unable to parse given date format'], ), ( - {'value': {Any(): 1, Any(): 2}}, - 'Index must be all timestamp like column. Valid formats are: pandas Timestamp, datetime, string in format %Y-%m-%d %H:%M:%S. Got .', + [Any(), Any()], + ['Got ."], ), ] @pytest.mark.parametrize('data', invalid_cases) def test_detect_and_parse_datetime_index_error_cases(mlflow_repository, data): - input_data = DataFrame(data[0]) + input_data = DataFrame({'value': [1, 2]}, index=data[0]) - message = data[1] + expected_fragments = data[1] with pytest.raises(ValueError) as e: mlflow_repository.detect_and_parse_datetime_index(input_data, metadata['metadata']) - assert str(e) == message + message = str(e.value) + + assert 'Index must be all timestamp like column' in message + assert f'string in format {DATETIME_FORMAT_WITH_TZ}' in message + for fragment in expected_fragments: + assert fragment in message, f'Expected {fragment!r} in {message!r}' valid_cases = [ @@ -1713,3 +1734,357 @@ async def test_predict_success_ndarray(mlflow_repository): assert 'prediction' in content.columns assert 'response_time' in content.columns assert content.index.tolist() == data.index.tolist() + + +""" +Tests for the joblib fallback on the model load paths +""" + +JOBLIB_PATH = 'laborious.utils.repository.model_repository.joblib' + + +def build_artifact_dir(tmp_path, *filenames): + """Create a fake downloaded artifact directory with the given top-level files.""" + for filename in filenames: + (tmp_path / filename).write_bytes(b'pickle-bytes') + return str(tmp_path) + + +def emitted_tags(emit_metric_mock, metric_object): + """Collect the `tags` kwarg of every emit_metric call for a given metric.""" + return [ + emitted.kwargs['tags'] + for emitted in emit_metric_mock.call_args_list + if emitted.kwargs['metric_object'] is metric_object + ] + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_flavor_success_skips_fallback( + joblib_mock, mlflow, mlflow_repository +): + result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') + + assert result == mlflow.sklearn.load_model.return_value + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_flavor_success_skips_fallback( + joblib_mock, mlflow, mlflow_repository +): + mlflow_repository.get_model_run_id = MagicMock() + mlflow_repository.get_model_uri = MagicMock() + + result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') + + assert result == mlflow.sklearn.load_model.return_value + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_download_model_wrapper_success_skips_fallback( + joblib_mock, mlflow, mlflow_repository +): + mlflow_repository.dowload_artifacts = AsyncMock() + + result = await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True) + + assert result == ( + mlflow.pyfunc.load_model.return_value._model_impl.python_model, + mlflow_repository.dowload_artifacts.return_value, + ) + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository): + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.load_predict_model('test_model', {}, 'invalid') + + assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + mlflow_repository.emit_metric.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_invalid_flavor_skips_fallback(joblib_mock, mlflow_repository): + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.load_transform_model('test_model', {}, 'invalid') + + assert str(exc_info.value) == INVALID_FLAVOR_MESSAGE + joblib_mock.load.assert_not_called() + mlflow_repository.client.download_artifacts.assert_not_called() + mlflow_repository.emit_metric.assert_not_called() + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_fallback_recovers_load( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'MLmodel', 'model.pkl' + ) + + result = await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') + + assert result is joblib_mock.load.return_value + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='test_model', stage='Production' + ) + mlflow_repository.client.download_artifacts.assert_called_once_with( + 'run-123', 'prediction_model' + ) + joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl')) + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_fallback_recovers_load( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.get_model_uri = MagicMock() + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'MLmodel', 'model.pkl' + ) + + result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') + + assert result is joblib_mock.load.return_value + mlflow_repository.get_model_run_id.assert_called_with( + model_name='test_model', stage='Production' + ) + mlflow_repository.client.download_artifacts.assert_called_once_with('run-123', 'data_model') + joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl')) + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_download_model_wrapper_fallback_reuses_downloaded_dir( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl') + mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir) + mlflow.pyfunc.load_model.side_effect = Exception('flavor boom') + + model, returned_path = await mlflow_repository.download_model( + 'test_model', {}, 'predict', 'pyfunc', True + ) + + mlflow_repository.dowload_artifacts.assert_called_once_with( + 'test_model', {}, 'prediction_model' + ) + mlflow_repository.client.download_artifacts.assert_not_called() + joblib_mock.load.assert_called_once_with(str(tmp_path / 'python_model.pkl')) + + assert model is joblib_mock.load.return_value + assert returned_path == artifact_dir + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_both_attempts_fail_chains_causes( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + flavor_error = Exception('flavor boom') + mlflow.sklearn.load_model.side_effect = flavor_error + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'model.pkl' + ) + joblib_mock.load.side_effect = ValueError('joblib boom') + + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') + + assert str(exc_info.value) == 'joblib boom' + assert exc_info.value.__cause__ is flavor_error + + error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) + assert len(error_tags) == 1 + assert error_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback' + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_both_attempts_fail_emits_error_metric( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.get_model_uri = MagicMock() + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'model.pkl' + ) + joblib_mock.load.side_effect = ValueError('joblib boom') + + with pytest.raises(ValueError): + await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') + + error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) + assert len(error_tags) == 1 + assert error_tags[0]['operation_type'] == 'load_transform_model_joblib_fallback' + assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == [] + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_predict_model_fallback_success_metrics( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'model.pkl' + ) + mlflow_repository.warning = MagicMock() + + await mlflow_repository.load_predict_model('test_model', {}, 'sklearn') + + read_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) + assert len(read_tags) == 1 + assert read_tags[0]['operation_type'] == 'load_predict_model_joblib_fallback' + assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) == [] + + mlflow_repository.observe_lag.assert_called_once() + lag_call = mlflow_repository.observe_lag.call_args + assert lag_call.args[1] is metrics.MODEL_READ_LAG + assert lag_call.args[2]['operation_type'] == 'load_predict_model_joblib_fallback' + + mlflow_repository.warning.assert_called_once() + assert 'flavor boom' in mlflow_repository.warning.call_args.args[0] + + +@pytest.mark.parametrize( + ('filenames', 'expected'), + [ + (('MLmodel', 'model.pkl'), 'model.pkl'), + (('MLmodel', 'python_model.pkl'), 'python_model.pkl'), + (('b_model.pkl', 'a_model.pkl'), 'a_model.pkl'), + ], +) +def test_find_model_pickle_selection(mlflow_repository, tmp_path, filenames, expected): + local_dir = build_artifact_dir(tmp_path, *filenames) + + result = mlflow_repository._find_model_pickle( + local_dir, 'prediction_model', Exception('flavor boom') + ) + + assert result == str(tmp_path / expected) + + +def test_find_model_pickle_ignores_nested_pickles(mlflow_repository, tmp_path): + nested = tmp_path / 'artifacts' + nested.mkdir() + (nested / 'stacking_model.pkl').write_bytes(b'pickle-bytes') + + with pytest.raises(ValueError) as exc_info: + mlflow_repository._find_model_pickle( + str(tmp_path), 'prediction_model', Exception('flavor boom') + ) + + assert 'No .pkl file found' in str(exc_info.value) + + +def test_find_model_pickle_empty_directory(mlflow_repository, tmp_path): + with pytest.raises(ValueError) as exc_info: + mlflow_repository._find_model_pickle(str(tmp_path), 'data_model', Exception('flavor boom')) + + assert 'data_model' in str(exc_info.value) + + +def test_find_model_pickle_error_message_carries_context(mlflow_repository, tmp_path): + local_dir = build_artifact_dir(tmp_path, 'MLmodel') + + with pytest.raises(ValueError) as exc_info: + mlflow_repository._find_model_pickle( + local_dir, 'prediction_model', Exception('flavor boom') + ) + + message = str(exc_info.value) + assert 'prediction_model' in message + assert 'flavor boom' in message + assert 'MLmodel' in message + + +def test_find_model_pickle_ignores_directory_named_pkl(mlflow_repository, tmp_path): + (tmp_path / 'algo.pkl').mkdir() + + with pytest.raises(ValueError) as exc_info: + mlflow_repository._find_model_pickle( + str(tmp_path), 'prediction_model', Exception('flavor boom') + ) + + assert 'algo.pkl' in str(exc_info.value) + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_get_model_caches_model_loaded_by_fallback( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'model.pkl' + ) + + first = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn') + second = await mlflow_repository.get_model('test_model', {}, 5, 'predict', 'sklearn') + + assert first is joblib_mock.load.return_value + assert second is first + assert mlflow.sklearn.load_model.call_count == 1 + assert joblib_mock.load.call_count == 1 + + +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_predict_reports_both_causes_when_fallback_fails( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(return_value='run-123') + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'model.pkl' + ) + joblib_mock.load.side_effect = ValueError('joblib boom') + + data = DataFrame({'feat_1': {'index_1': 2, 'index_2': 3}}) + + output = await mlflow_repository.predict( + 'test_model', + data, + {'retention_minutes': 0, 'predict_flavor': 'sklearn'}, + metadata['metadata'], + ) + + assert output['success'] is False + assert 'joblib boom' in output['content']['traceback'] + assert 'flavor boom' in output['content']['traceback'] + + +def test_load_model_pickle_real_joblib_round_trip(mlflow_repository, tmp_path): + estimator = LinearRegression().fit(np.array([[0.0], [1.0], [2.0]]), np.array([0.0, 1.0, 2.0])) + joblib_lib.dump(estimator, tmp_path / 'model.pkl', compress=3) + + loaded = mlflow_repository._load_model_pickle( + str(tmp_path), 'prediction_model', Exception('flavor boom') + ) + + assert loaded is not estimator + assert loaded.predict(np.array([[3.0]]))[0] == pytest.approx(3.0) From d7bf36605d40d8cd9aae6d137461f2bb2f0a79cc Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Wed, 19 Aug 2026 11:01:52 -0300 Subject: [PATCH 2/4] QTZPOC-13 QTZPOC-13: Dynamically resolve MLflow pyfunc model code paths to support varied layouts during re-logging. --- .../utils/repository/model_repository.py | 39 ++++++++++++++++- .../utils/repository/test_model_repository.py | 43 ++++++++++++++++--- 2 files changed, 75 insertions(+), 7 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index d6dc2bc..ab007d4 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -1094,6 +1094,43 @@ class MLFlowRepository(SientiaMonitoring): } return retrain_data + def resolve_code_paths(self, artifact_path: str | None, metadata: dict) -> list[str]: + """Resolve the code dependencies to forward when re-logging a pyfunc model. + + MLflow flattens whatever was given as `code_paths` into `/code/`, + so the layout differs per model: some ship a `utils` package, others a + single self-contained module (`_model.py`). Re-logging must forward + the entries that are actually on disk — assuming one fixed layout makes + `_validate_and_copy_code_paths` raise `FileNotFoundError` for the other. + + Args: + artifact_path (str | None): Local directory of the downloaded model + artifact, e.g. `./tmp/artifacts//prediction_model`. + metadata (dict): Metadata for structured logging. + + Returns: + list[str]: Existing paths inside `/code`, empty when + the model carries no code directory. + """ + if not artifact_path: + self.warning('No artifact path available, logging model without code paths', metadata) + return [] + + code_dir = path.join(artifact_path, 'code') + + if not path.isdir(code_dir): + self.warning( + f'No code directory found at {code_dir}, logging model without code paths', + metadata, + ) + return [] + + return [ + path.join(code_dir, entry) + for entry in sorted(listdir(code_dir)) + if entry != '__pycache__' + ] + async def log_model(self, model_data: dict, flavor: str, model_type: str, metadata: dict): """Log a model into the active MLflow run. @@ -1114,7 +1151,7 @@ class MLFlowRepository(SientiaMonitoring): if flavor == 'sklearn': mlflow.sklearn.log_model(model, model_type) elif flavor == 'pyfunc': - code_path = [path.join(model_data['artifact_path'], 'code', 'utils')] + code_path = self.resolve_code_paths(model_data.get('artifact_path'), metadata) self.debug(f'Code path: {code_path}', metadata) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 64c5e5f..457e29f 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1049,21 +1049,52 @@ async def test_log_model_sklearn(mlflow, mlflow_repository): ) -@patch('laborious.utils.repository.model_repository.path') +def test_resolve_code_paths_package_layout(tmp_path, mlflow_repository): + code_dir = tmp_path / 'code' + (code_dir / 'utils').mkdir(parents=True) + (code_dir / '__pycache__').mkdir() + + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [ + str(code_dir / 'utils') + ] + + +def test_resolve_code_paths_single_module_layout(tmp_path, mlflow_repository): + code_dir = tmp_path / 'code' + code_dir.mkdir() + (code_dir / 'loop_valve_classifier_model.py').touch() + + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [ + str(code_dir / 'loop_valve_classifier_model.py') + ] + + +def test_resolve_code_paths_without_code_dir(tmp_path, mlflow_repository): + assert mlflow_repository.resolve_code_paths(str(tmp_path), metadata['metadata']) == [] + + +def test_resolve_code_paths_without_artifact_path(mlflow_repository): + assert mlflow_repository.resolve_code_paths(None, metadata['metadata']) == [] + + @pytest.mark.asyncio -async def test_log_model_pyfunc(path, mlflow, mlflow_repository): +async def test_log_model_pyfunc(tmp_path, mlflow, mlflow_repository): + code_dir = tmp_path / 'code' + code_dir.mkdir() + (code_dir / 'loop_valve_classifier_model.py').touch() + model_mock = MagicMock() - model_data = {'model': model_mock, 'artifact_path': 'artifact_path'} + model_data = {'model': model_mock, 'artifact_path': str(tmp_path)} await mlflow_repository.log_model( model_data, 'pyfunc', 'prediction_model', metadata['metadata'] ) mlflow.pyfunc.log_model.assert_not_called() - path.join.assert_called_once_with('artifact_path', 'code', 'utils') - model_mock.store_model.assert_called_once_with( - artifact_path='prediction_model', code_path=[path.join.return_value], to_disk=False + artifact_path='prediction_model', + code_path=[str(code_dir / 'loop_valve_classifier_model.py')], + to_disk=False, ) mlflow_repository.observe_lag.assert_called_once_with(ANY, metrics.MODEL_WRITE_LAG, ANY) From 03ac4b31e521ec880fdc774ad7fbefc400104917 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 20 Aug 2026 08:28:45 -0300 Subject: [PATCH 3/4] QTZPOC-13 fix coverage --- .../utils/repository/test_model_repository.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 457e29f..5ea4ce7 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1972,6 +1972,30 @@ async def test_load_transform_model_both_attempts_fail_emits_error_metric( assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == [] +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_download_model_wrapper_both_attempts_fail_chains_causes( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + flavor_error = Exception('flavor boom') + artifact_dir = build_artifact_dir(tmp_path, 'MLmodel', 'python_model.pkl') + mlflow_repository.dowload_artifacts = AsyncMock(return_value=artifact_dir) + mlflow.pyfunc.load_model.side_effect = flavor_error + joblib_mock.load.side_effect = ValueError('joblib boom') + + with pytest.raises(ValueError) as exc_info: + await mlflow_repository.download_model('test_model', {}, 'predict', 'pyfunc', True) + + assert str(exc_info.value) == 'joblib boom' + assert exc_info.value.__cause__ is flavor_error + + error_tags = emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_ERROR_COUNT) + assert len(error_tags) == 1 + assert error_tags[0]['operation_type'] == 'download_model_wrapper_joblib_fallback' + assert emitted_tags(mlflow_repository.emit_metric, metrics.MODEL_READ_COUNT) == [] + mlflow_repository.observe_lag.assert_not_called() + + @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_load_predict_model_fallback_success_metrics( From 291d16efa53a4f05667b9dc17234048015bd4b67 Mon Sep 17 00:00:00 2001 From: vitor-aignosi Date: Thu, 20 Aug 2026 10:23:24 -0300 Subject: [PATCH 4/4] QTZPOC-13 QTZPOC-13: Centralize MLflow joblib fallback logic, ensuring consistent model version resolution by reusing the initial run ID, improving logging, and updating dependencies. --- .../utils/repository/model_repository.py | 235 ++++++++++++------ requirements-light.txt | 3 +- requirements.txt | 3 +- .../utils/repository/test_model_repository.py | 22 ++ 4 files changed, 179 insertions(+), 84 deletions(-) diff --git a/laborious/utils/repository/model_repository.py b/laborious/utils/repository/model_repository.py index ab007d4..e89b685 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,9 +16,11 @@ Capabilities: import ctypes import gc +import inspect import threading import time import traceback +from collections.abc import Callable from datetime import datetime, timedelta from io import StringIO from os import environ, listdir, makedirs, path @@ -393,7 +395,11 @@ class MLFlowRepository(SientiaMonitoring): return path.join(local_dir, pkl_files[0]) def _load_model_pickle( - self, local_dir: str, artifact_name: str, flavor_error: Exception + self, + local_dir: str, + artifact_name: str, + flavor_error: Exception, + metadata: dict[str, Any] | None = None, ) -> Any: """ Deserialize a downloaded model artifact with `joblib`. @@ -407,6 +413,9 @@ class MLFlowRepository(SientiaMonitoring): local_dir (str): Local directory holding the downloaded artifact. artifact_name (str): Artifact name, e.g. 'prediction_model'. flavor_error (Exception): Error raised by the MLflow flavor loader. + metadata (dict[str, Any] | None): Metadata used for structured + logging, so the chosen pickle path correlates with the other + fallback log lines. Returns: Any: The object deserialized from the artifact's pickle file. @@ -417,7 +426,7 @@ class MLFlowRepository(SientiaMonitoring): """ pkl_path = self._find_model_pickle(local_dir, artifact_name, flavor_error) - self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}') + self.debug(f'joblib fallback: loading {artifact_name} from {pkl_path}', metadata) return joblib.load(pkl_path) @@ -427,6 +436,7 @@ class MLFlowRepository(SientiaMonitoring): artifact_name: str, flavor_error: Exception, metadata: dict[str, Any], + run_id: str | None = None, ) -> Any: """ Download a model artifact and deserialize it with `joblib`. @@ -438,11 +448,20 @@ class MLFlowRepository(SientiaMonitoring): artifact is always a `.pkl` and the deserialized object is returned as is. + Callers that already resolved a run to build the flavor URI must pass it + as `run_id`: resolving the Production run again here would open a window + where a promotion between the two lookups makes the fallback serve a + different version than the one the flavor attempt tried and logged. Only + callers whose flavor attempt used a symbolic URI — `models://production`, + resolved by MLflow itself — leave it as `None`. + Args: model_name (str): Registered model name. artifact_name (str): Artifact to download, e.g. 'prediction_model'. flavor_error (Exception): Error raised by the MLflow flavor loader. metadata (dict[str, Any]): Metadata used for structured logging. + run_id (str | None): Run already resolved by the caller. When `None`, + the Production run is resolved here. Returns: Any: The object deserialized from the artifact's pickle file. @@ -452,7 +471,8 @@ class MLFlowRepository(SientiaMonitoring): Exception: Any error raised while resolving the Production run, downloading the artifact, or deserializing the pickle. """ - run_id = self.get_model_run_id(model_name=model_name, stage='Production') + if run_id is None: + run_id = self.get_model_run_id(model_name=model_name, stage='Production') self.info( f'joblib fallback: downloading {artifact_name} of {model_name} from run {run_id}', @@ -461,7 +481,80 @@ class MLFlowRepository(SientiaMonitoring): local_dir = self.client.download_artifacts(run_id, artifact_name) - return self._load_model_pickle(local_dir, artifact_name, flavor_error) + return self._load_model_pickle(local_dir, artifact_name, flavor_error, metadata) + + async def _load_with_joblib_fallback( + self, + primary_loader: Callable[[], Any], + fallback_loader: Callable[[Exception], Any], + failure_description: str, + metadata: dict[str, Any], + fallback_labels: dict[str, Any], + core_labels: dict[str, Any] | None = None, + ) -> Any: + """ + Run a model load with the joblib fallback and its observability. + + Single owner of the fallback flow shared by every load path: warn on the + flavor failure without emitting an error metric, try the fallback, and + emit lag/count under the fallback labels — or the error count chained to + the flavor error when the fallback also fails. Fixes to the flow belong + here, not in the call sites. + + Both loaders are called with no observability concern of their own: + `fallback_loader` receives the flavor error and may be sync or async, + which is what lets the download-based paths and the already-downloaded + wrapper path share this method. + + Args: + primary_loader (Callable[[], Any]): Flavor load attempt. Any + exception it raises triggers the fallback. + fallback_loader (Callable[[Exception], Any]): joblib fallback, + receiving the flavor error. Its result is awaited when + awaitable. + failure_description (str): What failed, e.g. 'Flavor loader sklearn + failed for prediction model x'. The fallback and original error + are appended to it. + metadata (dict[str, Any]): Metadata used for structured logging. + fallback_labels (dict[str, Any]): Metric labels for the fallback + path, already carrying its `operation_type`. + core_labels (dict[str, Any] | None): Metric labels for a successful + flavor load. `None` emits no metric on that path. + + Returns: + Any: The model loaded by the flavor loader or by the fallback. + + Raises: + Exception: The fallback error, chained from the flavor error, when + both attempts fail. + """ + start_time = time.time() + try: + model = primary_loader() + except Exception as flavor_error: + self.warning( + f'{failure_description}, trying joblib fallback. Original error: {flavor_error}', + metadata, + ) + try: + model = fallback_loader(flavor_error) + if inspect.isawaitable(model): + model = await model + except Exception as fallback_error: + await self.emit_metric( + metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels + ) + raise fallback_error from flavor_error + + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + return model + + if core_labels is not None: + await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) + await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) + + return model async def load_predict_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' @@ -492,37 +585,26 @@ class MLFlowRepository(SientiaMonitoring): fallback_labels = self.get_core_labels( metadata, operation_type='load_predict_model_joblib_fallback' ) - start_time = time.time() - try: + + def load_flavor() -> Any: if flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(model_uri) - elif flavor == 'sklearn': - model = mlflow.sklearn.load_model(model_uri) - else: - model = mlflow.pytorch.load_model(model_uri) - except Exception as flavor_error: - self.warning( - f'Flavor loader {flavor} failed for prediction model {model_name}, ' - f'trying joblib fallback. Original error: {flavor_error}', - metadata, - ) - try: - model = await self._download_and_load_pickle( - model_name, 'prediction_model', flavor_error, metadata - ) - except Exception as fallback_error: - await self.emit_metric( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels - ) - raise fallback_error from flavor_error + return mlflow.pyfunc.load_model(model_uri) + if flavor == 'sklearn': + return mlflow.sklearn.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) - await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) - return model - - await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) - return model + return await self._load_with_joblib_fallback( + primary_loader=load_flavor, + fallback_loader=lambda flavor_error: self._download_and_load_pickle( + model_name, 'prediction_model', flavor_error, metadata + ), + failure_description=( + f'Flavor loader {flavor} failed for prediction model {model_name}' + ), + metadata=metadata, + fallback_labels=fallback_labels, + core_labels=core_labels, + ) async def load_transform_model( self, model_name: str, metadata: dict[str, Any], flavor: str = 'sklearn' @@ -545,6 +627,11 @@ class MLFlowRepository(SientiaMonitoring): Raises: Exception: If the model run ID or URI cannot be retrieved, or if the model cannot be loaded. + + Notes: + - The Production run is resolved once and reused by the joblib + fallback, so a promotion between both steps cannot make the + fallback serve a different version than the flavor attempt. """ if flavor not in ('sklearn', 'pyfunc', 'pytorch'): @@ -559,38 +646,28 @@ class MLFlowRepository(SientiaMonitoring): fallback_labels = self.get_core_labels( metadata, operation_type='load_transform_model_joblib_fallback' ) - start_time = time.time() - try: + + def load_flavor() -> Any: if flavor == 'sklearn': - model = mlflow.sklearn.load_model(model_uri) - elif flavor == 'pyfunc': - model = mlflow.pyfunc.load_model(model_uri) - else: - model = mlflow.pytorch.load_model(model_uri) - except Exception as flavor_error: - self.warning( - f'Flavor loader {flavor} failed for data model {model_name}, ' - f'trying joblib fallback. Original error: {flavor_error}', + return mlflow.sklearn.load_model(model_uri) + if flavor == 'pyfunc': + return mlflow.pyfunc.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) + + return await self._load_with_joblib_fallback( + primary_loader=load_flavor, + fallback_loader=lambda flavor_error: self._download_and_load_pickle( + model_name, + 'data_model', + flavor_error, metadata, - ) - try: - model = await self._download_and_load_pickle( - model_name, 'data_model', flavor_error, metadata - ) - except Exception as fallback_error: - await self.emit_metric( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels - ) - raise fallback_error from flavor_error - - await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) - return model - - await self.observe_lag(start_time, metrics.MODEL_READ_LAG, core_labels) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=core_labels) - - return model + run_id=latest_production_id, + ), + failure_description=f'Flavor loader {flavor} failed for data model {model_name}', + metadata=metadata, + fallback_labels=fallback_labels, + core_labels=core_labels, + ) async def download_model( self, @@ -637,26 +714,20 @@ class MLFlowRepository(SientiaMonitoring): fallback_labels = self.get_core_labels( metadata, operation_type='download_model_wrapper_joblib_fallback' ) - start_time = time.time() - try: - raw_model = mlflow.pyfunc.load_model(artifact_path) - model = raw_model._model_impl.python_model - except Exception as flavor_error: - self.warning( - f'Wrapper loader failed for {model_type} model {model_name}, ' - f'trying joblib fallback. Original error: {flavor_error}', - metadata, - ) - try: - model = self._load_model_pickle(artifact_path, target, flavor_error) - except Exception as fallback_error: - await self.emit_metric( - metric_object=metrics.MODEL_READ_ERROR_COUNT, tags=fallback_labels - ) - raise fallback_error from flavor_error - await self.observe_lag(start_time, metrics.MODEL_READ_LAG, fallback_labels) - await self.emit_metric(metric_object=metrics.MODEL_READ_COUNT, tags=fallback_labels) + def load_wrapper_flavor() -> Any: + raw_model = mlflow.pyfunc.load_model(artifact_path) + return raw_model._model_impl.python_model + + model = await self._load_with_joblib_fallback( + primary_loader=load_wrapper_flavor, + fallback_loader=lambda flavor_error: self._load_model_pickle( + artifact_path, target, flavor_error, metadata + ), + failure_description=f'Wrapper loader failed for {model_type} model {model_name}', + metadata=metadata, + fallback_labels=fallback_labels, + ) self.debug( f'Model wrapper loaded: {model.__class__.__name__}:{model.__dict__}', metadata diff --git a/requirements-light.txt b/requirements-light.txt index d5128de..48cb400 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -15,5 +15,6 @@ hyperopt shap pycurl scipy<1.14.0 -scikit-learn==1.5.2 +scikit-learn==1.4.2 +pycaret==3.3.2 joblib diff --git a/requirements.txt b/requirements.txt index 8f994e9..4b2a2a8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,5 +15,6 @@ hyperopt shap pycurl scipy<1.14.0 -scikit-learn==1.5.2 +scikit-learn==1.4.2 +pycaret==3.3.2 joblib diff --git a/tests/laborious/utils/repository/test_model_repository.py b/tests/laborious/utils/repository/test_model_repository.py index 5ea4ce7..4a89d69 100644 --- a/tests/laborious/utils/repository/test_model_repository.py +++ b/tests/laborious/utils/repository/test_model_repository.py @@ -1903,6 +1903,28 @@ async def test_load_transform_model_fallback_recovers_load( joblib_mock.load.assert_called_once_with(str(tmp_path / 'model.pkl')) +@patch(JOBLIB_PATH) +@pytest.mark.asyncio +async def test_load_transform_model_fallback_reuses_resolved_run( + joblib_mock, mlflow, mlflow_repository, tmp_path +): + """A promotion between both steps must not change the run served.""" + mlflow.sklearn.load_model.side_effect = Exception('flavor boom') + mlflow_repository.get_model_run_id = MagicMock(side_effect=['run-A', 'run-B']) + mlflow_repository.get_model_uri = MagicMock() + mlflow_repository.client.download_artifacts.return_value = build_artifact_dir( + tmp_path, 'MLmodel', 'model.pkl' + ) + + result = await mlflow_repository.load_transform_model('test_model', {}, 'sklearn') + + assert result is joblib_mock.load.return_value + mlflow_repository.get_model_run_id.assert_called_once_with( + model_name='test_model', stage='Production' + ) + mlflow_repository.client.download_artifacts.assert_called_once_with('run-A', 'data_model') + + @patch(JOBLIB_PATH) @pytest.mark.asyncio async def test_download_model_wrapper_fallback_reuses_downloaded_dir(