[REFACTOR][CI] Modernize Python dependency management with uv (#18812)

mshr-h · web-flow · commit f8dc0fb46f5b · 2026-02-24T21:04:55.000-08:00
This PR modernizes Python tooling in CI by migrating Python interpreter
installation, venv creation and package installation to uv.
diff --git a/docker/Dockerfile.ci_arm b/docker/Dockerfile.ci_arm
@@ -19,6 +19,7 @@
 # tag: v0.02
 
 FROM ubuntu:22.04
+COPY --from=ghcr.io/astral-sh/uv:0.10.5 /uv /uvx /bin/
 
 COPY utils/apt-install-and-clear.sh /usr/local/bin/apt-install-and-clear
 COPY utils/download-and-verify.sh /usr/local/bin/download-and-verify
@@ -47,11 +48,15 @@ ENV PATH=/opt/sccache:$PATH
 COPY install/ubuntu_install_llvm.sh /install/ubuntu_install_llvm.sh
 RUN bash /install/ubuntu_install_llvm.sh
 
+ENV UV_NO_CACHE=1
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_MANAGED_PYTHON=1
+
 ENV TVM_VENV=/venv/apache-tvm-py3.10
-COPY python/bootstrap/lockfiles /install/python/bootstrap/lockfiles
 COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
 RUN bash /install/ubuntu_install_python.sh 3.10
 ENV PATH=${TVM_VENV}/bin:$PATH
+ENV VIRTUAL_ENV=${TVM_VENV}
 ENV PYTHONNOUSERSITE=1
 
 COPY install/ubuntu_install_python_package.sh /install/ubuntu_install_python_package.sh
diff --git a/docker/Dockerfile.ci_cpu b/docker/Dockerfile.ci_cpu
@@ -17,6 +17,7 @@
 
 # CI docker CPU env
 FROM ubuntu:22.04
+COPY --from=ghcr.io/astral-sh/uv:0.10.5 /uv /uvx /bin/
 
 COPY utils/apt-install-and-clear.sh /usr/local/bin/apt-install-and-clear
 COPY utils/download-and-verify.sh /usr/local/bin/download-and-verify
@@ -35,11 +36,16 @@ RUN bash /install/ubuntu_install_cmake.sh
 COPY install/ubuntu_install_googletest.sh /install/ubuntu_install_googletest.sh
 RUN bash /install/ubuntu_install_googletest.sh
 
+ENV UV_NO_CACHE=1
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_MANAGED_PYTHON=1
+
 ENV TVM_VENV=/venv/apache-tvm-py3.10
-COPY python/bootstrap/lockfiles /install/python/bootstrap/lockfiles
 COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
 RUN bash /install/ubuntu_install_python.sh 3.10
+RUN uv --version
 ENV PATH=${TVM_VENV}/bin:$PATH
+ENV VIRTUAL_ENV=${TVM_VENV}
 
 # Disable .local directory from affecting CI.
 ENV PYTHONNOUSERSITE=1
diff --git a/docker/Dockerfile.ci_gpu b/docker/Dockerfile.ci_gpu
@@ -18,6 +18,7 @@
 # CI docker GPU env
 # tag: v0.60
 FROM nvidia/cuda:12.8.1-cudnn-devel-ubuntu22.04
+COPY --from=ghcr.io/astral-sh/uv:0.10.5 /uv /uvx /bin/
 
 COPY utils/apt-install-and-clear.sh /usr/local/bin/apt-install-and-clear
 COPY utils/download-and-verify.sh /usr/local/bin/download-and-verify
@@ -42,18 +43,20 @@ RUN bash /install/ubuntu_install_cmake.sh
 COPY install/ubuntu_install_googletest.sh /install/ubuntu_install_googletest.sh
 RUN bash /install/ubuntu_install_googletest.sh /googletest
 
+ENV UV_NO_CACHE=1
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_MANAGED_PYTHON=1
+
 ENV TVM_VENV=/venv/apache-tvm-py3.10
-COPY python/bootstrap/lockfiles /install/python/bootstrap/lockfiles
 COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
 RUN bash /install/ubuntu_install_python.sh 3.10
+RUN uv --version
 ENV PATH=${TVM_VENV}/bin:$PATH
+ENV VIRTUAL_ENV=${TVM_VENV}
 
 # Disable .local directory from affecting CI.
 ENV PYTHONNOUSERSITE=1
 
-# Globally disable pip cache
-RUN pip config set global.no-cache-dir false
-
 COPY install/ubuntu_install_llvm.sh /install/ubuntu_install_llvm.sh
 RUN bash /install/ubuntu_install_llvm.sh
 
diff --git a/docker/Dockerfile.ci_lint b/docker/Dockerfile.ci_lint
@@ -19,6 +19,7 @@
 # CI docker lint env
 # tag: v0.60
 FROM ubuntu:22.04
+COPY --from=ghcr.io/astral-sh/uv:0.10.5 /uv /uvx /bin/
 
 COPY utils/apt-install-and-clear.sh /usr/local/bin/apt-install-and-clear
 COPY utils/download-and-verify.sh /usr/local/bin/download-and-verify
@@ -31,16 +32,20 @@ RUN bash /install/ubuntu_setup_tz.sh
 RUN apt-get update && \
     apt-install-and-clear -y wget git sudo make parallel doxygen graphviz curl shellcheck ca-certificates curl gnupg
 
+ENV UV_NO_CACHE=1
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_MANAGED_PYTHON=1
+
 ENV TVM_VENV=/venv/apache-tvm-py3.10
-COPY python/bootstrap/lockfiles /install/python/bootstrap/lockfiles
 COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
 RUN bash /install/ubuntu_install_python.sh 3.10
 ENV PATH=${TVM_VENV}/bin:$PATH
+ENV VIRTUAL_ENV=${TVM_VENV}
 
 # Disable .local directory from affecting CI.
 ENV PYTHONNOUSERSITE=1
 
-RUN pip3 install cpplint==1.6.1 pylint==2.17.2 mypy==0.902 black==22.12.0 flake8==3.9.2 blocklint==0.2.3 jinja2==3.0.3
+RUN uv pip install cpplint==1.6.1 pylint==2.17.2 mypy==0.902 black==22.12.0 flake8==3.9.2 blocklint==0.2.3 jinja2==3.0.3
 
 # java deps for rat
 COPY install/ubuntu_install_java.sh /install/ubuntu_install_java.sh
diff --git a/docker/Dockerfile.ci_wasm b/docker/Dockerfile.ci_wasm
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 FROM ubuntu:22.04
+COPY --from=ghcr.io/astral-sh/uv:0.10.5 /uv /uvx /bin/
 
 COPY utils/apt-install-and-clear.sh /usr/local/bin/apt-install-and-clear
 COPY utils/download-and-verify.sh /usr/local/bin/download-and-verify
@@ -33,11 +34,15 @@ RUN bash /install/ubuntu_install_cmake.sh
 COPY install/ubuntu_install_googletest.sh /install/ubuntu_install_googletest.sh
 RUN bash /install/ubuntu_install_googletest.sh
 
+ENV UV_NO_CACHE=1
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+ENV UV_MANAGED_PYTHON=1
+
 ENV TVM_VENV=/venv/apache-tvm-py3.10
-COPY python/bootstrap/lockfiles /install/python/bootstrap/lockfiles
 COPY install/ubuntu_install_python.sh /install/ubuntu_install_python.sh
 RUN bash /install/ubuntu_install_python.sh 3.10
 ENV PATH=${TVM_VENV}/bin:$PATH
+ENV VIRTUAL_ENV=${TVM_VENV}
 
 # Disable .local directory from affecting CI.
 ENV PYTHONNOUSERSITE=1
diff --git a/docker/install/ubuntu_install_core.sh b/docker/install/ubuntu_install_core.sh
@@ -30,6 +30,7 @@ apt-get update && apt-install-and-clear -y --no-install-recommends \
     g++ \
     gdb \
     git \
+    gnupg \
     graphviz \
     libcurl4-openssl-dev \
     libopenblas-dev \
diff --git a/docker/install/ubuntu_install_coreml.sh b/docker/install/ubuntu_install_coreml.sh
@@ -20,4 +20,4 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install coremltools
+uv pip install coremltools==8.3.0
diff --git a/docker/install/ubuntu_install_cuda_python.sh b/docker/install/ubuntu_install_cuda_python.sh
@@ -20,4 +20,4 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install cuda-python
+uv pip install cuda-python==12.8.0
diff --git a/docker/install/ubuntu_install_dgl.sh b/docker/install/ubuntu_install_dgl.sh
@@ -20,4 +20,4 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install dgl==v1.0.0 -f https://data.dgl.ai/wheels/repo.html
+uv pip install dgl==v1.0.0 -f https://data.dgl.ai/wheels/repo.html
diff --git a/docker/install/ubuntu_install_gluoncv.sh b/docker/install/ubuntu_install_gluoncv.sh
@@ -20,4 +20,4 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install gluoncv
+uv pip install gluoncv==0.10.5.post0
diff --git a/docker/install/ubuntu_install_jax.sh b/docker/install/ubuntu_install_jax.sh
@@ -22,14 +22,14 @@ set -o pipefail
 
 # Install jax and jaxlib
 if [ "$1" == "cuda" ]; then
-    pip3 install --upgrade \
+    uv pip install --upgrade \
         jaxlib~=0.4.9 \
         "jax[cuda11_pip]~=0.4.9" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 else
-    pip3 install --upgrade \
+    uv pip install --upgrade \
         jaxlib~=0.4.9 \
         "jax[cpu]~=0.4.9"
 fi
 
 # Install flax
-pip3 install flax~=0.6.9
+uv pip install flax~=0.6.9
diff --git a/docker/install/ubuntu_install_nnef.sh b/docker/install/ubuntu_install_nnef.sh
@@ -20,5 +20,5 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install \
+uv pip install \
     nnef==1.0.8
diff --git a/docker/install/ubuntu_install_onnx.sh b/docker/install/ubuntu_install_onnx.sh
@@ -27,17 +27,17 @@ PYTHON_VERSION=$(python3 -c "import sys; print(f'{sys.version_info.major}.{sys.v
 DEVICE=${1:-cpu}
 
 # Install the onnx package
-pip3 install \
+uv pip install \
     onnx==1.20.1 \
     onnxruntime==1.23.2 \
     onnxoptimizer==0.4.2
 
 if [ "$DEVICE" == "cuda" ]; then
-    pip3 install \
+    uv pip install \
         torch==2.10.0 \
         torchvision==0.25.0
 else
-    pip3 install \
+    uv pip install \
         torch==2.10.0 \
         torchvision==0.25.0 \
         --extra-index-url https://download.pytorch.org/whl/cpu
diff --git a/docker/install/ubuntu_install_python.sh b/docker/install/ubuntu_install_python.sh
@@ -33,53 +33,26 @@ if [ "$#" -lt 1 ]; then
 fi
 PYTHON_VERSION=$1
 
-case "$PYTHON_VERSION" in
-    3.10|3.11) ;;
-    *)
-        echo "Only 3.10 and 3.11 versions are supported in this script."
-        exit -1
-        ;;
-esac
-
+# Install base dependencies required by this script.
 apt-get update
+apt-install-and-clear -y acl
 
-# Ensure lsb-release is installed.
-apt-install-and-clear -y lsb-core software-properties-common
-
-release=$(lsb_release -sc)
-case "${release}" in
-    jammy)
-        if [ "${PYTHON_VERSION}" == "3.10" ] || \
-           [ "${PYTHON_VERSION}" == "3.11" ]; then
-            add-apt-repository -y ppa:deadsnakes/ppa
-        fi
-        ;;
-    *)
-        echo "Don't know which version of python to install for lsb-release ${release}"
-        exit 2
-        ;;
-esac
-
-# Install python and pip. Don't modify this to add Python package dependencies,
-# instead modify install_python_package.sh
-apt-install-and-clear -y \
-    acl \
-    python${PYTHON_VERSION} \
-    python${PYTHON_VERSION}-dev \
-    python3-pip \
-    python${PYTHON_VERSION}-venv
-
-update-alternatives --install /usr/bin/python3 python3 /usr/bin/python${PYTHON_VERSION} 1
-update-alternatives --install /usr/bin/python python /usr/bin/python${PYTHON_VERSION} 1
+# Install managed Python to a shared location so both root and the CI runtime user can use it.
+UV_PYTHON_INSTALL_DIR=${UV_PYTHON_INSTALL_DIR:-/opt/uv/python}
+export UV_PYTHON_INSTALL_DIR
+export UV_MANAGED_PYTHON=${UV_MANAGED_PYTHON:-1}
+mkdir -p "${UV_PYTHON_INSTALL_DIR}"
+chmod 755 "${UV_PYTHON_INSTALL_DIR}"
+uv python install "${PYTHON_VERSION}"
 
 # Allow disabling user site-packages, even with sudo; this makes it harder to repro CI failures
 # locally because it's hard to tell what might be in this directory.
 echo "Defaults env_keep += \"PYTHONNOUSERSITE\"" >/etc/sudoers.d/91-preserve-python-nousersite
 export PYTHONNOUSERSITE=1
 
-venv_dir="$(python3 -c "import os.path;print(os.path.dirname(\"${TVM_VENV}\"))")"
+venv_dir="$(dirname "${TVM_VENV}")"
 mkdir -p "${venv_dir}"
-python3 -mvenv "${TVM_VENV}"
+uv venv --python "${PYTHON_VERSION}" "${TVM_VENV}"
 
 # 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.
@@ -88,32 +61,16 @@ 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.
-if [ -f "/install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt" ]; then
-    pip_spec=$(tac /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt | grep -m 1 'pip==')
-    ${TVM_VENV}/bin/pip install -U --require-hashes -r <(echo "${pip_spec}") \
-        -c /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
-
-    # Python configuration
-    ${TVM_VENV}/bin/pip config set global.no-cache-dir true  # Never cache packages
-
-    # Now install the remaining base packages.
-    ${TVM_VENV}/bin/pip install \
-        --require-hashes \
-        -r /install/python/bootstrap/lockfiles/constraints-${PYTHON_VERSION}.txt
-fi
-
-addgroup tvm-venv
+addgroup tvm-venv || true
 chgrp -R tvm-venv "${TVM_VENV}"
 setfacl -R -d -m group:tvm-venv:rwx "${TVM_VENV}"
 setfacl -R -m group:tvm-venv:rwx "${TVM_VENV}"
 
 # Prevent further use of pip3 via the system.
 # There may be multiple (i.e. from python3-pip apt package and pip3 install -U).
-while [ "$(which pip3)" != "" ]; do
-    rm "$(which pip3)"
+while command -v pip3 >/dev/null 2>&1; do
+    rm -f "$(command -v pip3)"
 done
-while [ "$(which pip)" != "" ]; do
-    rm "$(which pip)"
+while command -v pip >/dev/null 2>&1; do
+    rm -f "$(command -v pip)"
 done
diff --git a/docker/install/ubuntu_install_python_package.sh b/docker/install/ubuntu_install_python_package.sh
@@ -21,25 +21,25 @@ set -u
 set -o pipefail
 
 # install libraries for python package on ubuntu
-pip3 install --upgrade \
-    "Pygments>=2.4.0" \
-    cloudpickle \
-    cython \
-    mypy \
+uv pip install --upgrade \
+    "Pygments~=2.19" \
+    "cloudpickle~=3.1" \
+    "cython~=3.0" \
+    "mypy~=1.15" \
     numpy==1.26.* \
-    orderedset \
-    packaging \
+    "orderedset~=2.0" \
+    "packaging~=25.0" \
     Pillow==12.1.1 \
-    psutil \
-    pytest \
+    "psutil~=7.0" \
+    "pytest~=8.3" \
     git+https://github.com/tlc-pack/tlcpack-sphinx-addon.git@768ec1dce349fe4708f6ad68be1ebb3f3dabafa1 \
-    pytest-profiling \
-    pytest-xdist \
+    "pytest-profiling~=1.8" \
+    "pytest-xdist~=3.6" \
     pytest-rerunfailures==16.1 \
-    requests \
-    scipy \
-    Jinja2 \
+    "requests~=2.32" \
+    "scipy~=1.13" \
+    "Jinja2~=3.1" \
     junitparser==4.0.2 \
-    six \
-    tornado \
-    "ml_dtypes>=0.5.1"
+    "six~=1.17" \
+    "tornado~=6.4" \
+    "ml_dtypes~=0.5"
diff --git a/docker/install/ubuntu_install_redis.sh b/docker/install/ubuntu_install_redis.sh
@@ -21,5 +21,5 @@ set -u
 set -o pipefail
 
 apt-get update && apt-install-and-clear -y redis-server
-pip3 install \
+uv pip install \
     xgboost==3.2.0
diff --git a/docker/install/ubuntu_install_sphinx.sh b/docker/install/ubuntu_install_sphinx.sh
@@ -20,14 +20,14 @@ set -e
 set -u
 set -o pipefail
 
-pip3 install \
-    autodocsumm \
-    commonmark \
-    docutils \
-    Image \
-    matplotlib \
-    sphinx \
-    sphinx_autodoc_annotation \
-    sphinx-gallery \
-    sphinx_rtd_theme \
-    https://github.com/tlc-pack/tlcpack-sphinx-addon/archive/refs/tags/v0.2.3.zip
+uv pip install \
+    autodocsumm==0.2.7 \
+    commonmark==0.9.1 \
+    docutils==0.16 \
+    image==1.5.33 \
+    matplotlib==3.3.4 \
+    sphinx==4.2.0 \
+    sphinx_autodoc_annotation~=1.0 \
+    git+https://github.com/sphinx-gallery/sphinx-gallery.git@6142f1791151849b5bec4bf3959f75697ba226cd \
+    sphinx_rtd_theme==1.0.0 \
+    git+https://github.com/tlc-pack/tlcpack-sphinx-addon.git@768ec1dce349fe4708f6ad68be1ebb3f3dabafa1
diff --git a/docker/install/ubuntu_install_tensorflow.sh b/docker/install/ubuntu_install_tensorflow.sh
diff --git a/docker/install/ubuntu_install_tensorflow_aarch64.sh b/docker/install/ubuntu_install_tensorflow_aarch64.sh
diff --git a/docker/install/ubuntu_install_tflite.sh b/docker/install/ubuntu_install_tflite.sh
