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..e89b685 100644 --- a/laborious/utils/repository/model_repository.py +++ b/laborious/utils/repository/model_repository.py @@ -16,15 +16,18 @@ 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, 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 +354,208 @@ 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, + metadata: dict[str, Any] | None = None, + ) -> 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. + 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. + + 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}', metadata) + + 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], + run_id: str | None = None, + ) -> 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. + + 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. + + 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. + """ + 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}', + metadata, + ) + + local_dir = self.client.download_artifacts(run_id, artifact_name) + + 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' ) -> Any: @@ -370,27 +575,36 @@ 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') - 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 + fallback_labels = self.get_core_labels( + metadata, operation_type='load_predict_model_joblib_fallback' + ) - 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 + def load_flavor() -> Any: + if flavor == 'pyfunc': + return mlflow.pyfunc.load_model(model_uri) + if flavor == 'sklearn': + return mlflow.sklearn.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, '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' @@ -413,32 +627,47 @@ 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'): + 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') - start_time = time.time() - try: + fallback_labels = self.get_core_labels( + metadata, operation_type='load_transform_model_joblib_fallback' + ) + + def load_flavor() -> Any: 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 + return mlflow.sklearn.load_model(model_uri) + if flavor == 'pyfunc': + return mlflow.pyfunc.load_model(model_uri) + return mlflow.pytorch.load_model(model_uri) - 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, + 'data_model', + flavor_error, + metadata, + 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, @@ -482,8 +711,23 @@ 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' + ) + + 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 @@ -921,6 +1165,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. @@ -941,7 +1222,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/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 5daca26..f2481d4 100644 --- a/requirements-light.txt +++ b/requirements-light.txt @@ -15,4 +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-local.txt b/requirements-local.txt index 06a69d7..aac0487 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 c44b8a9..1f0fbf4 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +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 54a75bb..4a89d69 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 = [ @@ -1028,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) @@ -1713,3 +1765,403 @@ 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_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( + 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_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( + 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)