Update mlflow version in sdk

python/sdk/requirements.test.txt

@@ -0,0 +1,18 @@
+google-cloud-bigquery-storage>=0.7.0
+google-cloud-bigquery>=1.18.0
+joblib>=0.13.0,<1.2.0"  # >=1.2.0 upon upgrade of kserve's version
+mypy>=0.812
+pytest-cov
+pytest-dependency
+pytest-xdist
+pytest
+recursive-diff>=1.0.0
+requests
+scikit-learn>=1.1.2
+types-python-dateutil
+types-PyYAML
+types-six
+types-protobuf
+urllib3-mock>=0.3.3
+xarray
+xgboost==1.6.2

python/sdk/requirements.txt

@@ -0,0 +1,20 @@
+boto3>=1.35.39
+caraml-upi-protos>=0.3.1
+certifi>=2017.4.17
+Click>=7.0,<8.1.4
+cloudpickle==2.0.0  # used by mlflow
+cookiecutter>=1.7.2
+dataclasses-json>=0.5.2  # allow Flyte version 1.2.0 or above to import Merlin SDK
+docker<=6.1.3
+GitPython>=3.1.40
+google-cloud-storage>=1.19.0
+protobuf>=3.12.0,<5.0.0  # Determined by the mlflow dependency
+mlflow==2.16.2
+PyPrind>=2.11.2
+python_dateutil>=2.5.3
+PyYAML>=5.4
+six>=1.10
+urllib3>=1.26
+numpy<=1.23.5  # Temporary pin numpy due to https://numpy.org/doc/stable/release/1.20.0-notes.html#numpy-1-20-0-release-notes
+caraml-auth-google==0.0.0.post7
+pydantic==2.5.3

python/sdk/setup.py

@@ -15,56 +15,26 @@
 
 import imp
 import os
+import pathlib
+import pkg_resources
 
 from setuptools import find_packages, setup
 
 version = imp.load_source(
     "merlin.version", os.path.join("merlin", "version.py")
 ).VERSION
 
-REQUIRES = [
-    "boto3>=1.9.84",
-    "caraml-upi-protos>=0.3.1",
-    "certifi>=2017.4.17",
-    "Click>=7.0,<8.1.4",
-    "cloudpickle==2.0.0",  # used by mlflow
-    "cookiecutter>=1.7.2",
-    "dataclasses-json>=0.5.2",  # allow Flyte version 1.2.0 or above to import Merlin SDK
-    "docker<=6.1.3",
-    "GitPython>=3.1.40",
-    "google-cloud-storage>=1.19.0",
-    "protobuf>=3.12.0,<5.0.0",  # Determined by the mlflow dependency
-    "mlflow==1.26.1",
-    "PyPrind>=2.11.2",
-    "python_dateutil>=2.5.3",
-    "PyYAML>=5.4",
-    "six>=1.10",
-    "urllib3>=1.26",
-    "numpy<=1.23.5",  # Temporary pin numpy due to https://numpy.org/doc/stable/release/1.20.0-notes.html#numpy-1-20-0-release-notes
-    "caraml-auth-google==0.0.0.post7",
-    "pydantic==2.5.3"
-]
+with pathlib.Path("requirements.txt").open() as requirements_txt:
+    requirements = [
+        str(requirement)
+        for requirement in pkg_resources.parse_requirements(requirements_txt)
+    ]
 
-TEST_REQUIRES = [
-    "google-cloud-bigquery-storage>=0.7.0",
-    "google-cloud-bigquery>=1.18.0",
-    "joblib>=0.13.0,<1.2.0",  # >=1.2.0 upon upgrade of kserve's version
-    "mypy>=0.812",
-    "pytest-cov",
-    "pytest-dependency",
-    "pytest-xdist",
-    "pytest",
-    "recursive-diff>=1.0.0",
-    "requests",
-    "scikit-learn>=1.1.2",
-    "types-python-dateutil",
-    "types-PyYAML",
-    "types-six",
-    "types-protobuf",
-    "urllib3-mock>=0.3.3",
-    "xarray",
-    "xgboost==1.6.2",
-]
+with pathlib.Path("requirements.test.txt").open() as test_requirements_test:
+    test_requirements = [
+        str(requirement)
+        for requirement in pkg_resources.parse_requirements(test_requirements_test)
+    ]
 
 setup(
     name="merlin-sdk",
@@ -75,10 +45,10 @@
     packages=find_packages(),
     package_data={"merlin": ["docker/pyfunc.Dockerfile", "docker/standard.Dockerfile"]},
     zip_safe=True,
-    install_requires=REQUIRES,
+    install_requires=requirements,
     setup_requires=["setuptools_scm"],
-    tests_require=TEST_REQUIRES,
-    extras_require={"test": TEST_REQUIRES},
+    tests_require=test_requirements,
+    extras_require={"test": test_requirements},
     python_requires=">=3.8,<3.11",
     long_description=open("README.md").read(),
     long_description_content_type="text/markdown",

python/sdk/test/pyfunc_test.py

@@ -55,7 +55,7 @@ def test_http():
     pyfunc_model = PyFuncModel()
     pyfunc_model.infer = MagicMock(return_value=response)
 
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {
@@ -71,7 +71,7 @@ def test_http_headers():
     pyfunc_model = PyFuncModel()
     pyfunc_model.infer = MagicMock(return_value=response)
 
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {
@@ -88,7 +88,7 @@ def test_explicit_protocol():
     pyfunc_model = PyFuncModel()
     pyfunc_model.infer = MagicMock(return_value=response)
 
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {
@@ -105,7 +105,7 @@ def test_upiv1():
     pyfunc_model = PyFuncModel()
     pyfunc_model.upiv1_infer = MagicMock(return_value=upiv1_response)
 
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {
@@ -128,7 +128,7 @@ def test_pyfuncv3_rest():
     pyfunc_model.preprocess = MagicMock(return_value=pyfunc_model_input)
     pyfunc_model.infer = MagicMock(return_value=pyfunc_model_output)
     pyfunc_model.postprocess = MagicMock(return_value=http_response)
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {
@@ -187,7 +187,7 @@ def test_pyfuncv3_upi():
     pyfunc_model.upiv1_preprocess = MagicMock(return_value=pyfunc_model_input)
     pyfunc_model.infer = MagicMock(return_value=pyfunc_model_output)
     pyfunc_model.upiv1_postprocess = MagicMock(return_value=upi_response)
-    context = PythonModelContext(artifacts={})
+    context = PythonModelContext(artifacts={}, model_config={})
 
     pyfunc_model.load_context(context)
     model_input = {