Dockerfile.base: PyPI versions of pip[-tools]

Author: olupton

.github/container/Dockerfile.base

@@ -17,6 +17,9 @@ ARG CLANG_VERSION
 ARG JAX_TOOLBOX_REF
 ENV CUDA_BASE_IMAGE=${BASE_IMAGE}
 
+# Install all python packages system-wide.
+ENV PIP_BREAK_SYSTEM_PACKAGES=1
+
 ###############################################################################
 ## Install Python and essential tools
 ###############################################################################
@@ -38,6 +41,7 @@ apt_packages=(
   liblzma-dev
   python-is-python3
   python3-pip
+  python3.12-dev
   rsync
   vim
   wget
@@ -46,8 +50,6 @@ apt_packages=(
   # llvm.sh
   lsb-release 
   software-properties-common
-  # GCP autoconfig
-  pciutils hwloc bind9-host
 )
 if [[ $(dpkg --print-architecture) == arm64 ]]; then
   # h5py: The newest release of of h5py (3.11.0) does not include ARM wheels and causes pip to build h5py.
@@ -90,12 +92,16 @@ EOL
 apt-get clean
 rm -rf /var/lib/apt/lists/*
 
+# Upgrade pip from PyPI before we remove the Ubuntu version
+pip install --ignore-installed --no-cache-dir --upgrade pip
+
 # There are several python packages (in the list below) that are installed with OS 
 # package manager (the run of `apt-get install` above) and can not be uninstall
 # using pip (in pip-finalize.sh script) during JAX installation. Remove then in
 # advance to avoid JAX installation issue.
 remove_packages=(
   python3-gi
+  python3-pip
   software-properties-common
   lsb-release
   python3-yaml
@@ -117,24 +123,9 @@ ADD --chmod=777 \
   /usr/local/bin/
 RUN wget https://github.com/mikefarah/yq/releases/latest/download/yq_linux_$(dpkg --print-architecture) -O /usr/local/bin/yq && \
   chmod 777 /usr/local/bin/yq
-RUN git clone -b 23.3.1 https://github.com/pypa/pip.git /opt/pip
-# Patch is specific to 23.3.1
-# Generated via: "git diff > pip-vcs-equivalency.patch"
-ADD pip-vcs-equivalency.patch /opt/pip/
-RUN <<EOF bash -e -x
-cd /opt/pip
-git apply </opt/pip/pip-vcs-equivalency.patch
-git add -u
-git commit -m 'Adds JAX_TOOLBOX_VCS_EQUIVALENCY as a trigger to treat all github VCS installs for a package as equivalent. The spec of the last encountered version will be used'
-EOF
 
-# install all python packages system-wide.
-ENV PIP_BREAK_SYSTEM_PACKAGES=1
-# An extra flag `--ignore-installed` is added below, because of the following reason: 
-# after upgrading to ver 23.3.1 (from /opt/pip) `pip` tries to uninstall itself (default pip-24.0) 
-# and fails due to pip-24.0 has been installed with system tool `apt` but not `python`. So we keep 
-# both pip-24.0 and pip-23.3.1 in the system, but use 23.3.1 with equivalency patch (see above).
-RUN pip install --upgrade --ignore-installed --no-cache-dir -e /opt/pip pip-tools && rm -rf ~/.cache/*
+# pip-tools are used by pip-finalize.sh
+RUN pip install --no-cache-dir pip-tools
 
 ###############################################################################
 ## Symlink for cuDNN

.github/container/pip-finalize.sh

@@ -28,15 +28,8 @@ done | tee requirements.vcs
 unset IFS
 
 # Second pip-compile includes one more requirements file that pins all vcs installs
-# Uses a special env var to let our custom pip impl know to treat the following as
-# equivalent:
-#
-# fiddle @ git+https://github.com/google/fiddle
-# fiddle @ git+https://github.com/google/fiddle@cd4497e4c09bdf95dcccaa1e138c2c125d32d39f
-#
-# JAX_TOOLBOX_VCS_EQUIVALENCY is an environment variable enabling custom logic in pip
-# that treats the above as equivalent and prefers the URI wit the SHA
-JAX_TOOLBOX_VCS_EQUIVALENCY=true pip-compile -o requirements.txt requirements.vcs $(ls requirements-*.in)
+# NB there was an old patch to do with equivalency; kw JAX_TOOLBOX_VCS_EQUIVALENCY
+pip-compile -o requirements.txt requirements.vcs $(ls requirements-*.in)
 
 # If there are unpinned VCS dependencies, error since these should be included in the manifest
 unpinned_vcs_dependencies=$(cat requirements.txt | egrep '^[^#].+ @ git\+' | egrep -v '^[^#].+ @ git\+.+@' || true)