feat: update mlflow version to support newer version of protobuf

caraml-dev · Oct 3, 2023 · 7e312c1 · 7e312c1
1 parent 072325c
commit 7e312c1
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 14 deletions.
diff --git a/python/batch-predictor/docker/app.Dockerfile b/python/batch-predictor/docker/app.Dockerfile
@@ -24,6 +24,6 @@ ARG GOOGLE_APPLICATION_CREDENTIALS
 RUN if [[ ! -z "$GOOGLE_APPLICATION_CREDENTIALS" ]]; then gcloud auth activate-service-account --key-file=${GOOGLE_APPLICATION_CREDENTIALS}; fi
 RUN gsutil -m cp -r ${MODEL_URL} .
 RUN /bin/bash -c ". activate ${CONDA_ENVIRONMENT} && \
-    sed -i 's/mlflow$/mlflow==1.23.0/' ${HOME}/model/conda.yaml && \
+    sed -i 's/\bmlflow[^[:space:]]*/mlflow==1.26.1/' ${HOME}/model/conda.yaml && \
     conda env update --name ${CONDA_ENVIRONMENT} --file ${HOME}/model/conda.yaml && \
     python ${HOME}/merlin-spark-app/main.py --dry-run-model ${HOME}/model"
diff --git a/python/batch-predictor/merlinpyspark/model.py b/python/batch-predictor/merlinpyspark/model.py
@@ -67,7 +67,7 @@ def spark_udf(spark, model_uri, features, result_type="double"):
         archive_path = SparkModelCache.add_local_model(spark, local_model_path)
 
     def predict(*args):
-        model = SparkModelCache.get_or_load(archive_path)
+        model, _ = SparkModelCache.get_or_load(archive_path)
         schema = {features[i]: arg for i, arg in enumerate(args)}
         pdf = None
         for x in args:

diff --git a/python/batch-predictor/requirements.txt b/python/batch-predictor/requirements.txt
@@ -1,8 +1,6 @@
 findspark
 pyspark==3.0.1
-mlflow>=1.2.0,<=1.23.0
+mlflow>=1.26.1,<=2.0.0
 cloudpickle==2.0.0
 pyarrow>=0.14.1,<=9.0.0
-protobuf>=3.0,<4.0.0
-#TODO: Update merlin-sdk dep to: file:${SDK_PATH}#egg=merlin-sdk
-merlin-sdk==0.33.0 # Pin to the version that supports Python 3.7.
+file:${SDK_PATH}#egg=merlin-sdk
diff --git a/python/batch-predictor/setup.py b/python/batch-predictor/setup.py
@@ -21,12 +21,11 @@
 with open('requirements.txt') as f:
     REQUIRE = f.read().splitlines()
 
-# TODO: Uncomment below lines after Pyfunc server stops supporting Python 3.7
-# merlin_path = os.path.join(os.getcwd(), "../sdk")
-# merlin_sdk_package = "merlin-sdk"
-# for index, item in enumerate(REQUIRE):
-#     if merlin_sdk_package in item:
-#         REQUIRE[index] = f"{merlin_sdk_package} @ file://localhost/{merlin_path}#egg={merlin_sdk_package}"
+merlin_path = os.path.join(os.getcwd(), "../sdk")
+merlin_sdk_package = "merlin-sdk"
+for index, item in enumerate(REQUIRE):
+    if merlin_sdk_package in item:
+        REQUIRE[index] = f"{merlin_sdk_package} @ file://localhost/{merlin_path}#egg={merlin_sdk_package}"
 
 setup(
     name='merlin-pyspark-app',

diff --git a/python/sdk/setup.py b/python/sdk/setup.py
@@ -30,8 +30,7 @@
     "cookiecutter>=1.7.2",
     "docker>=4.2.1",
     "google-cloud-storage>=1.19.0",
-    "mlflow>=1.2.0,<=1.23.0", #  for py3.11 due to proto -> "mlflow>=1.26.1",
-    "protobuf>=3.0.0,<4.0.0", #  for py3.11 due to proto -> "protobuf>=4.0.0,<5.0dev",
+    "mlflow>=1.26.1,<2.0.0",
     "PyPrind>=2.11.2",
     "python_dateutil>=2.5.3",
     "PyYAML>=5.4",