Merge pull request #39699 from apache/cp-39683

[Cherrypick] fix ensurepip bundled pip cleanup for Python 3.12+ containers

Author: Amar3tto

sdks/python/container/Dockerfile

@@ -49,8 +49,6 @@ RUN  \
 
     rm -rf /var/lib/apt/lists/* && \
 
-    pip install --upgrade pip setuptools wheel && \
-
     # Install required packages for Beam Python SDK and common dependencies used by users.
     # use --no-deps to ensure the list includes all transitive dependencies.
     # use --prefer-binary to avoid compiling wheels from source when prebuilt wheels exist.
@@ -92,17 +90,13 @@ RUN  \
     if [ "${py_version}" = "3.10" ] || [ "${py_version}" = "3.11" ]; then \
         pip install upgrade_ensurepip; \
         python3 -m upgrade_ensurepip; \
-    else \
-        python3 /tmp/upgrade_bundled_pip.py; \
-    fi; \
-    # setuptools is not bundled with ensurepip in Python 3.12+
-    if [ "${py_version}" = "3.10" ] || [ "${py_version}" = "3.11" ]; then \
         find /usr/local/lib/python${py_version}/ensurepip/_bundled/setuptools-* -type f ! -name $(basename $(ls -v /usr/local/lib/python${py_version}/ensurepip/_bundled/setuptools-*-py3-none-any.whl | tail -n 1)) -delete; \
-    fi; \
-    find /usr/local/lib/python${py_version}/ensurepip/_bundled/pip-* -type f ! -name $(basename $(ls -v /usr/local/lib/python${py_version}/ensurepip/_bundled/pip-*-py3-none-any.whl | tail -n 1)) -delete; \
-    if [ "${py_version}" = "3.10" ] || [ "${py_version}" = "3.11" ]; then \
+        find /usr/local/lib/python${py_version}/ensurepip/_bundled/pip-* -type f ! -name $(basename $(ls -v /usr/local/lib/python${py_version}/ensurepip/_bundled/pip-*-py3-none-any.whl | tail -n 1)) -delete; \
         pip uninstall upgrade_ensurepip -y; \
+    else \
+        python3 /tmp/upgrade_bundled_pip.py; \
     fi; \
+    # Verify ensurepip can bootstrap pip. Required by boot.go worker venv creation.
     python3 -m ensurepip;
 
 ENTRYPOINT ["/opt/apache/beam/boot"]

sdks/python/container/license_scripts/upgrade_bundled_pip.py

@@ -21,6 +21,10 @@
 The script is executed within Docker after the image pip has been upgraded.
 upgrade_ensurepip expects setuptools to be bundled as well, but Python 3.12+
 only ships pip in ensurepip/_bundled.
+
+After downloading, removes other pip-* files from _bundled so only the wheel
+matching the installed pip version remains. Worker harness startup (boot.go)
+creates a venv via ensurepip, so that wheel must be present.
 """
 
 import subprocess
@@ -34,8 +38,8 @@ def main():
   ep_path = Path(ensurepip.__file__)
   wheel_dir = ep_path.parent / '_bundled'
   pip_version = subprocess.check_output(
-      [sys.executable, '-m', 'pip', '--version'],
-      text=True).split()[1]
+      [sys.executable, '-m', 'pip', '--version'], text=True).split()[1]
+  expected_wheel_name = 'pip-{}-py3-none-any.whl'.format(pip_version)
   subprocess.check_call([
       sys.executable,
       '-m',
@@ -46,6 +50,13 @@ def main():
       str(wheel_dir),
       '--no-deps',
   ])
+  for path in wheel_dir.glob('pip-*'):
+    if path.name != expected_wheel_name:
+      path.unlink()
+  if not (wheel_dir / expected_wheel_name).is_file():
+    sys.exit(
+        'ensurepip bundled pip wheel missing after install: {}'.format(
+            wheel_dir / expected_wheel_name))
   lines = ep_path.read_text().splitlines()
   pip_line = None
   for idx, line in enumerate(lines):