Blunt approach with --break-system-packages

NVIDIA · Oct 28, 2024 · 8ed3daf · 8ed3daf
1 parent ab0c427
commit 8ed3daf
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/.github/container/Dockerfile.base b/.github/container/Dockerfile.base
@@ -8,17 +8,17 @@ ARG CLANG_VERSION=18
 ## Obtain GCP's NCCL TCPx plugin
 ###############################################################################
 
-FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 as tcpx-installer-amd64
+FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 AS tcpx-installer-amd64
 
 # make a stub arm64 container because GCP does not provide an arm64 version of the plugin
-FROM ubuntu as tcpx-installer-arm64
+FROM ubuntu AS tcpx-installer-arm64
 RUN <<"OUTEREOF" bash -ex
 mkdir -p /scripts /var/lib/tcpx/lib64
 echo '#!/bin/bash' > /scripts/container_entry.sh
 chmod +x /scripts/container_entry.sh
 OUTEREOF
 
-FROM tcpx-installer-${TARGETARCH} as tcpx-installer
+FROM tcpx-installer-${TARGETARCH} AS tcpx-installer
 RUN /scripts/container_entry.sh install
 
 ###############################################################################
@@ -50,7 +50,6 @@ apt_packages=(
   gnupg
   liblzma-dev
   python-is-python3
-  python3-pip
   rsync
   vim
   wget
@@ -125,7 +124,24 @@ 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
-RUN pip install --upgrade --no-cache-dir --break-system-packages -e /opt/pip pip-tools && rm -rf ~/.cache/*
+RUN <<EOF bash -e
+PIP_PYZ_DIR=$(mktemp -d)
+pushd "${PIP_PYZ_DIR}"
+curl -O https://bootstrap.pypa.io/pip/pip.pyz
+chmod +x ./pip.pyz
+! pip && ! pip3 # there should not already be a pip
+python pip.pyz install --break-system-packages --no-cache-dir -e /opt/pip
+popd
+rm -rf "${PIP_PYZ_DIR}"
+# Make sure that any later attempt to install `python3-pip` will fail
+cat >/etc/apt/preferences.d/no-system-pip <<EOL
+# A patched pip is installed at /opt/pip.
+Package: python3-pip
+Pin: release *
+Pin-Priority: -1
+EOL
+pip install --break-system-packages --no-cache-dir pip-tools
+EOF
 
 ###############################################################################
 ## Install TCPx