From 76d919f811150a63bbdfa5edc631fb77a5562752 Mon Sep 17 00:00:00 2001
From: Krithika Sundararajan <krithika.sundararajan@gojek.com>
Date: Wed, 29 Nov 2023 14:26:10 +0800
Subject: [PATCH] Use SDK version that matches PyFunc servers

---
 python/batch-predictor/docker/app.Dockerfile | 2 +-
 python/pyfunc-server/docker/Dockerfile       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/batch-predictor/docker/app.Dockerfile b/python/batch-predictor/docker/app.Dockerfile
index 5d6d7de23..2d1002e0f 100644
--- a/python/batch-predictor/docker/app.Dockerfile
+++ b/python/batch-predictor/docker/app.Dockerfile
@@ -25,6 +25,6 @@ RUN if [[ ! -z "$GOOGLE_APPLICATION_CREDENTIALS" ]]; then gcloud auth activate-s
 RUN gsutil -m cp -r ${MODEL_URL} .
 RUN /bin/bash -c ". activate ${CONDA_ENVIRONMENT} && \
     sed -i 's/mlflow[><=]\{0,1\}.*$/mlflow==1.26.1/g' ${HOME}/model/conda.yaml && \
-    sed -i '/merlin-sdk/d' ${HOME}/model/conda.yaml && \
     conda env update --name ${CONDA_ENVIRONMENT} --file ${HOME}/model/conda.yaml && \
+    pip install ${SDK_PATH} && \
     python ${HOME}/merlin-spark-app/main.py --dry-run-model ${HOME}/model"
diff --git a/python/pyfunc-server/docker/Dockerfile b/python/pyfunc-server/docker/Dockerfile
index b33a19e58..bde7e32d7 100644
--- a/python/pyfunc-server/docker/Dockerfile
+++ b/python/pyfunc-server/docker/Dockerfile
@@ -25,8 +25,8 @@ RUN if [ ! -z "$GOOGLE_APPLICATION_CREDENTIALS" ]; then gcloud auth activate-ser
 RUN gsutil cp -r ${MODEL_URL} .
 RUN /bin/bash -c ". activate merlin-model && \
     sed -i 's/mlflow[><=]\{0,1\}.*$/mlflow==1.26.1/g' model/conda.yaml && \
-    sed -i '/merlin-sdk/d' model/conda.yaml && \
     conda env update --name merlin-model --file model/conda.yaml && \
+    pip install ${SDK_PATH} && \
     python -m pyfuncserver --model_dir model --dry_run"
 
 CMD ["/bin/bash", "./run.sh"]