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:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user