@@ -8,17 +8,17 @@ ARG CLANG_VERSION=18
88## Obtain GCP's NCCL TCPx plugin
99###############################################################################
1010
11- FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 as tcpx-installer-amd64
11+ FROM us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx:v3.1.10 AS tcpx-installer-amd64
1212
1313# make a stub arm64 container because GCP does not provide an arm64 version of the plugin
14- FROM ubuntu as tcpx-installer-arm64
14+ FROM ubuntu AS tcpx-installer-arm64
1515RUN <<"OUTEREOF" bash -ex
1616mkdir -p /scripts /var/lib/tcpx/lib64
1717echo '#!/bin/bash' > /scripts/container_entry.sh
1818chmod +x /scripts/container_entry.sh
1919OUTEREOF
2020
21- FROM tcpx-installer-${TARGETARCH} as tcpx-installer
21+ FROM tcpx-installer-${TARGETARCH} AS tcpx-installer
2222RUN /scripts/container_entry.sh install
2323
2424###############################################################################
@@ -50,7 +50,6 @@ apt_packages=(
5050 gnupg
5151 liblzma-dev
5252 python-is-python3
53- python3-pip
5453 rsync
5554 vim
5655 wget
@@ -125,7 +124,24 @@ git apply </opt/pip/pip-vcs-equivalency.patch
125124git add -u
126125git 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'
127126EOF
128- RUN pip install --upgrade --no-cache-dir --break-system-packages -e /opt/pip pip-tools && rm -rf ~/.cache/*
127+ RUN <<EOF bash -e
128+ PIP_PYZ_DIR=$(mktemp -d)
129+ pushd "${PIP_PYZ_DIR}"
130+ curl -O https://bootstrap.pypa.io/pip/pip.pyz
131+ chmod +x ./pip.pyz
132+ ! pip && ! pip3 # there should not already be a pip
133+ python pip.pyz install --break-system-packages --no-cache-dir -e /opt/pip
134+ popd
135+ rm -rf "${PIP_PYZ_DIR}"
136+ # Make sure that any later attempt to install `python3-pip` will fail
137+ cat >/etc/apt/preferences.d/no-system-pip <<EOL
138+ # A patched pip is installed at /opt/pip.
139+ Package: python3-pip
140+ Pin: release *
141+ Pin-Priority: -1
142+ EOL
143+ pip install --break-system-packages --no-cache-dir pip-tools
144+ EOF
129145
130146###############################################################################
131147## Install TCPx
0 commit comments