Reinstall onnx and tf deps when python version changes (#1652)

Author: vishalbollu

dev/versions.md

@@ -179,7 +179,7 @@ Note: it's ok if example training notebooks aren't upgraded, as long as the expo
 ## Python packages
 
 1. Update versions in `images/python-predictor-*/Dockerfile`, `images/tensorflow-predictor/Dockerfile`, and `images/onnx-predictor-*/Dockerfile`
-1. Update versions in `pkg/workloads/cortex/serve/requirements.txt` and `pkg/workloads/cortex/downloader/requirements.txt`
+1. Update versions in `pkg/workloads/cortex/serve/*requirements.txt` and `pkg/workloads/cortex/downloader/requirements.txt`
 1. Update the versions listed in "Pre-installed packages" in `realtime-api/predictors.md` and `batch-api/predictors.md`
     * look at the diff carefully since some packages are not shown, and e.g. `tensorflow-cpu` -> `tensorflow`
     * be careful not to update any of the versions for Inferentia that are not latest in `images/python-predictor-inf/Dockerfile`

images/onnx-predictor-cpu/Dockerfile

@@ -44,9 +44,11 @@ ENV BASH_ENV=~/.env
 SHELL ["/bin/bash", "-c"]
 
 COPY pkg/workloads/cortex/serve/requirements.txt /src/cortex/serve/requirements.txt
+COPY pkg/workloads/cortex/serve/onnx-cpu.requirements.txt /src/cortex/serve/image.requirements.txt
+
 RUN pip install --no-cache-dir -r \
     /src/cortex/serve/requirements.txt \
-    onnxruntime==1.4.0
+    -r /src/cortex/serve/image.requirements.txt
 
 ARG SLIM="false"
 RUN test "${SLIM}" = "true" || ( \

images/onnx-predictor-gpu/Dockerfile

@@ -44,9 +44,11 @@ ENV BASH_ENV=~/.env
 SHELL ["/bin/bash", "-c"]
 
 COPY pkg/workloads/cortex/serve/requirements.txt /src/cortex/serve/requirements.txt
+COPY pkg/workloads/cortex/serve/onnx-gpu.requirements.txt /src/cortex/serve/image.requirements.txt
+
 RUN pip install --no-cache-dir -r \
     /src/cortex/serve/requirements.txt \
-    onnxruntime-gpu==1.4.0
+    -r /src/cortex/serve/image.requirements.txt
 
 ARG SLIM="false"
 RUN test "${SLIM}" = "true" || ( \

images/tensorflow-predictor/Dockerfile

@@ -44,10 +44,11 @@ ENV BASH_ENV=~/.env
 SHELL ["/bin/bash", "-c"]
 
 COPY pkg/workloads/cortex/serve/requirements.txt /src/cortex/serve/requirements.txt
+COPY pkg/workloads/cortex/serve/tf.requirements.txt /src/cortex/serve/image.requirements.txt
+
 RUN pip install --no-cache-dir -r \
     /src/cortex/serve/requirements.txt \
-    tensorflow-cpu==2.3.0 \
-    tensorflow-serving-api==2.3.0
+    -r /src/cortex/serve/image.requirements.txt
 
 ARG SLIM="false"
 RUN test "${SLIM}" = "true" || ( \

pkg/workloads/cortex/serve/init/bootloader.sh

@@ -74,7 +74,11 @@ if [ -f "/mnt/project/conda-packages.txt" ]; then
     if [ $old_py_version != $new_py_version ]; then
         echo "warning: you have changed the Python version from $old_py_version to $new_py_version; this may break Cortex's web server"
         echo "reinstalling core packages ..."
+
         pip --no-cache-dir install -r /src/cortex/serve/requirements.txt
+        if [ -f "/src/cortex/serve/image.requirements.txt" ]; then
+            pip --no-cache-dir install -r /src/cortex/serve/image.requirements.txt
+        fi
 
         rm -rf $CONDA_PREFIX/lib/python${old_py_version}  # previous python is no longer needed
     fi

pkg/workloads/cortex/serve/onnx-cpu.requirements.txt

@@ -0,0 +1 @@
+onnxruntime==1.4.0

pkg/workloads/cortex/serve/onnx-gpu.requirements.txt

@@ -0,0 +1 @@
+onnxruntime-gpu==1.4.0

pkg/workloads/cortex/serve/tf.requirements.txt

@@ -0,0 +1,2 @@
+tensorflow-cpu==2.3.0
+tensorflow-serving-api==2.3.0