chore: import upstream snapshot with attribution

tests/pyfunc/test_pyfunc_model_config.py

@@ -0,0 +1,154 @@
+import os
+
+import pytest
+import yaml
+
+import mlflow
+from mlflow.models import Model
+
+
+@pytest.fixture
+def model_path(tmp_path):
+    return os.path.join(tmp_path, "model")
+
+
+@pytest.fixture
+def model_config():
+    return {
+        "use_gpu": True,
+        "temperature": 0.9,
+        "timeout": 300,
+    }
+
+
+def _load_pyfunc(path):
+    return TestModel()
+
+
+class TestModel(mlflow.pyfunc.PythonModel):
+    def predict(self, context, model_input, params=None):
+        return model_input
+
+
+class InferenceContextModel(mlflow.pyfunc.PythonModel):
+    def predict(self, context, model_input, params=None):
+        # This mock class returns the internal inference configuration keys and values available
+        return context.model_config.items()
+
+
+def test_save_with_model_config(model_path, model_config):
+    model = InferenceContextModel()
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config)
+
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
+
+    assert loaded_model.model_config
+    assert set(model_config.keys()) == set(loaded_model.model_config)
+    assert all(loaded_model.model_config[k] == v for k, v in model_config.items())
+    assert all(loaded_model.model_config[k] == v for k, v in loaded_model.predict([[0]]))
+
+
+@pytest.mark.parametrize(
+    "model_config_path",
+    [
+        os.path.abspath("tests/pyfunc/sample_code/config.yml"),
+        "tests/pyfunc/../pyfunc/sample_code/config.yml",
+    ],
+)
+def test_save_with_model_config_path(model_path, model_config, model_config_path):
+    model = InferenceContextModel()
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config_path)
+
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path)
+
+    assert loaded_model.model_config
+    assert set(model_config.keys()) == set(loaded_model.model_config)
+    assert all(loaded_model.model_config[k] == v for k, v in model_config.items())
+    assert all(loaded_model.model_config[k] == v for k, v in loaded_model.predict([[0]]))
+
+
+def test_override_model_config(model_path, model_config):
+    model = TestModel()
+    inference_override = {"timeout": 400}
+
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config)
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path, model_config=inference_override)
+
+    assert loaded_model.model_config["timeout"] == 400
+    assert all(loaded_model.model_config[k] == v for k, v in inference_override.items())
+
+
+@pytest.mark.parametrize(
+    "model_config_path",
+    [
+        os.path.abspath("tests/pyfunc/sample_code/config.yml"),
+        "tests/pyfunc/../pyfunc/sample_code/config.yml",
+    ],
+)
+def test_override_model_config_path(tmp_path, model_path, model_config_path):
+    model = TestModel()
+    inference_override = {"timeout": 400}
+    config_path = tmp_path / "config.yml"
+    config_path.write_text(yaml.dump(inference_override))
+
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config_path)
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path, model_config=str(config_path))
+
+    assert loaded_model.model_config["timeout"] == 400
+    assert all(loaded_model.model_config[k] == v for k, v in inference_override.items())
+
+
+def test_override_model_config_ignore_invalid(model_path, model_config):
+    model = TestModel()
+    inference_override = {"invalid_key": 400}
+
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config)
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path, model_config=inference_override)
+
+    assert loaded_model.predict([[5]])
+    assert all(k not in loaded_model.model_config for k in inference_override.keys())
+
+
+@pytest.mark.parametrize(
+    "model_config_path",
+    [
+        os.path.abspath("tests/pyfunc/sample_code/config.yml"),
+        "tests/pyfunc/../pyfunc/sample_code/config.yml",
+    ],
+)
+def test_override_model_config_path_ignore_invalid(tmp_path, model_path, model_config_path):
+    model = TestModel()
+    inference_override = {"invalid_key": 400}
+    config_path = tmp_path / "config.yml"
+    config_path.write_text(yaml.dump(inference_override))
+
+    mlflow.pyfunc.save_model(model_path, python_model=model, model_config=model_config_path)
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path, model_config=str(config_path))
+
+    assert loaded_model.predict([[5]])
+    assert all(k not in loaded_model.model_config for k in inference_override.keys())
+
+
+def test_pyfunc_without_model_config(model_path, model_config):
+    model = TestModel()
+    mlflow.pyfunc.save_model(model_path, python_model=model)
+
+    loaded_model = mlflow.pyfunc.load_model(model_uri=model_path, model_config=model_config)
+
+    assert loaded_model.predict([[5]])
+    assert not loaded_model.model_config
+
+
+def test_pyfunc_loader_without_model_config(model_path):
+    mlflow.pyfunc.save_model(
+        path=model_path,
+        data_path=".",
+        loader_module=__name__,
+        code_paths=[__file__],
+        mlflow_model=Model(run_id="test", artifact_path="testtest"),
+    )
+
+    inference_override = {"invalid_key": 400}
+    pyfunc_model = mlflow.pyfunc.load_model(model_path, model_config=inference_override)
+
+    assert not pyfunc_model.model_config