gitea_admin/sientia-dataops-laborious-legacy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into fix/SIENTIAPDE-2072

Browse Source
This commit is contained in:
Vitor Pimentel dos Santos
2026-08-25 16:32:39 -03:00
committed by GitHub
parent 695e4c07a6 4576502d78
commit 83b025ff9f
17 changed files with 1279 additions and 204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@@ -55,4 +55,5 @@ mlruns/
relatorio*
openspec/
.cursor/
.cursor/
.claude/

242
e2e/conftest.py
View File

@@ -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.

19
e2e/helpers.py
View File

@@ -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(

7
e2e/opc_test_server.py
View File

@@ -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)

4
e2e/test_child_workflows_e2e.py
View File

@@ -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': {

24
e2e/test_minio_offload.py
View File

@@ -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 = {

29
e2e/test_opc_real_server.py
View File

@@ -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(

98
e2e/test_predictions_batch_format_export.py
View File

@@ -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.",
)

186
e2e/test_predictions_batch_joblib_fallback.py Normal file
View File

@@ -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

3
e2e/test_predictions_batch_main_workflow.py
View File

@@ -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

14
e2e/test_predictions_batch_prediction_process.py
View File

@@ -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(

353
laborious/utils/repository/model_repository.py
View File

@@ -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:/<name>/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 `<artifact>/code/`,
so the layout differs per model: some ship a `utils` package, others a
single self-contained module (`<model>_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/<model>/prediction_model`.
metadata (dict): Metadata for structured logging.
Returns:
list[str]: Existing paths inside `<artifact_path>/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)

5
pyproject.toml
View File

@@ -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

4
requirements-light.txt
View File

@@ -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

4
requirements-local.txt
View File

@@ -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

4
requirements.txt
View File

@@ -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

484
tests/laborious/utils/repository/test_model_repository.py
View File

@@ -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 <class 'str'>, <class 'int'>.",
['2024-01-01 12:00:00', 2024],
["Elements are <class 'str'>,<class 'int'>"],
),
(
{'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 <class Any>.',
[Any(), Any()],
['Got <class', "Any'>."],
),
]
@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)