SAM0123-320: fix: call pyfunc data model predict with wrapper signature

When a data model's `fit` returns the refit object instead of a DataFrame,
`fit_models` called `predict(data)` sklearn-style. A pyfunc wrapper's
`predict` takes `(context, model_input)`, so Courier Stratified retrain
failed with "missing 1 required positional argument: 'model_input'".

Mirror `get_prediction_data`: pass `{}` as context when the transform
flavor is pyfunc.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
2026-10-07 16:47:31 -03:00
parent 25f11fa943
commit fdeafe9768
2 changed files with 36 additions and 1 deletions

View File

@@ -1341,6 +1341,9 @@ class MLFlowRepository(SientiaMonitoring):
if not isinstance(treated_data_candidate, pd.DataFrame):
data_model = treated_data_candidate
if transform_flavor == 'pyfunc':
treated_data = data_model.predict({}, data)
else:
treated_data = data_model.predict(data)
else:
treated_data = treated_data_candidate

View File

@@ -1043,6 +1043,38 @@ async def test_fit_models_not_df_target_name_none_and_not_in_model(
}
@patch('laborious.utils.repository.model_repository.pd.merge')
@patch('laborious.utils.repository.model_repository.isinstance')
@pytest.mark.asyncio
async def test_fit_models_pyfunc_transform_predicts_with_wrapper_signature(
isinstance_mock, pd_merge, mlflow_repository
):
"""A pyfunc data model refit returns the wrapper, whose `predict` takes `(context, model_input)`."""
isinstance_mock.return_value = False
data_model = MagicMock()
mlflow_repository.download_model = AsyncMock(
side_effect=[(data_model, 'artifact_path'), (MagicMock(), 'artifact_path')],
)
mlflow_repository.detect_and_parse_datetime_index = MagicMock()
mlflow_repository.get_prediction_data = MagicMock(return_value=DataFrame())
data = MagicMock()
await mlflow_repository.fit_models(
'model_name',
data,
'latest_production_id',
metadata['metadata'],
'pyfunc',
False,
'pyfunc',
'target',
)
data_model.fit.return_value.predict.assert_called_once_with({}, data)
@patch('laborious.utils.repository.model_repository.pd.merge')
@patch('laborious.utils.repository.model_repository.isinstance')
@pytest.mark.asyncio