From 7c59216fdfbee8eaf8b98866af9dc6b89838554f Mon Sep 17 00:00:00 2001
From: Andrew Reusch <areusch@gmail.com>
Date: Fri, 16 Sep 2022 16:54:44 -0700
Subject: [PATCH 1/2] Fix caffe, boost install in Python venvs by creating
 python3.X link.

---
 docker/install/ubuntu_install_boost.sh  | 3 +++
 docker/install/ubuntu_install_caffe.sh  | 3 +--
 docker/install/ubuntu_install_python.sh | 7 +++++++
 3 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/docker/install/ubuntu_install_boost.sh b/docker/install/ubuntu_install_boost.sh
index e226bbc5d96e..1c04c8a8ff20 100755
--- a/docker/install/ubuntu_install_boost.sh
+++ b/docker/install/ubuntu_install_boost.sh
@@ -24,6 +24,9 @@ cleanup() {
 
 trap cleanup 0
 
+# NOTE: by default, tvm-venv python is used. Install boost on the system.
+PATH=${PATH/${TVM_VENV}\/bin:/}
+
 curl -LO https://boostorg.jfrog.io/artifactory/main/release/1.67.0/source/boost_1_67_0.tar.gz
 BOOST_HASH=8c247e040303a97895cee9c9407ef205e2c3ab09f0b8320997835ad6221dff23a87231629498ccfd0acca473f74e9ec27b8bd774707b062228df1e5f72d44c92
 echo "$BOOST_HASH" boost_1_67_0.tar.gz | sha512sum -c
diff --git a/docker/install/ubuntu_install_caffe.sh b/docker/install/ubuntu_install_caffe.sh
index 4d9763b69aa3..af9fff6e26ac 100755
--- a/docker/install/ubuntu_install_caffe.sh
+++ b/docker/install/ubuntu_install_caffe.sh
@@ -65,5 +65,4 @@ cd / && rm -rf /caffe_src
 
 PYCAFFE_ROOT=${CAFFE_HOME}/python
 echo "${CAFFE_HOME}/lib" >> /etc/ld.so.conf.d/caffe.conf && ldconfig
-VENV_SITE_PACKAGE=$(pip3 show numpy | grep "Location:" | cut -d ' ' -f 2)
-ln -s ${PYCAFFE_ROOT}/caffe ${VENV_SITE_PACKAGE}/caffe
+ln -s ${PYCAFFE_ROOT}/caffe "${TVM_VENV}/usr/local/lib/python3.7/dist-packages/caffe"
diff --git a/docker/install/ubuntu_install_python.sh b/docker/install/ubuntu_install_python.sh
index 66a80e1fdc52..fb31c41dccea 100755
--- a/docker/install/ubuntu_install_python.sh
+++ b/docker/install/ubuntu_install_python.sh
@@ -65,6 +65,13 @@ mkdir -p "${venv_dir}"
 python3 -mvenv "${TVM_VENV}"
 . "${TVM_VENV}/bin/activate"
 
+# NOTE: Only in python3.9 does venv guarantee it creates the python3.X binary.
+# This is needed so that cmake's find_package(PythonInterp) works inside the venv.
+# See https://bugs.python.org/issue39656
+if [ ! -e "${TVM_VENV}/bin/python${PYTHON_VERSION}" ]; then
+    ln -s "${TVM_VENV}/bin/python" "${TVM_VENV}/bin/python${PYTHON_VERSION}"
+fi
+
 # Update pip to match version used to produce requirements-hashed.txt. This step
 # is necessary so that pip's dependency solver is recent.
 pip_spec=$(cat /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt | grep 'pip==')

From d0757bceda5fa4dfeec886523539b9afd84dba92 Mon Sep 17 00:00:00 2001
From: Andrew Reusch <areusch@gmail.com>
Date: Tue, 20 Sep 2022 11:15:05 -0700
Subject: [PATCH 2/2] Use getsitepackages()

---
 docker/install/ubuntu_install_caffe.sh | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/docker/install/ubuntu_install_caffe.sh b/docker/install/ubuntu_install_caffe.sh
index af9fff6e26ac..1e42270e267a 100755
--- a/docker/install/ubuntu_install_caffe.sh
+++ b/docker/install/ubuntu_install_caffe.sh
@@ -65,4 +65,5 @@ cd / && rm -rf /caffe_src
 
 PYCAFFE_ROOT=${CAFFE_HOME}/python
 echo "${CAFFE_HOME}/lib" >> /etc/ld.so.conf.d/caffe.conf && ldconfig
-ln -s ${PYCAFFE_ROOT}/caffe "${TVM_VENV}/usr/local/lib/python3.7/dist-packages/caffe"
+site_packages=$("${TVM_VENV}/bin/python3" -c 'import site; print(site.getsitepackages()[0])')
+ln -s ${PYCAFFE_ROOT}/caffe "${site_packages}/caffe"