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.
This commit is contained in:
242
e2e/conftest.py
242
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.
|
||||
|
||||
Reference in New Issue
Block a user