From ef988a4d139515f395dd857f5660610dec24d470 Mon Sep 17 00:00:00 2001
From: aibrahiim <abdoibrahim1017@gmail.com>
Date: Sat, 8 Aug 2026 00:00:21 +0300
Subject: [PATCH] fix ensurepip bundled pip cleanup for Python 3.12+ containers

---
 sdks/python/container/Dockerfile                  |  4 +++-
 .../license_scripts/upgrade_bundled_pip.py        | 15 +++++++++++++--
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/sdks/python/container/Dockerfile b/sdks/python/container/Dockerfile
index acb637521c92..8c6dd74a2f91 100644
--- a/sdks/python/container/Dockerfile
+++ b/sdks/python/container/Dockerfile
@@ -91,11 +91,13 @@ RUN  \
         pip install upgrade_ensurepip; \
         python3 -m upgrade_ensurepip; \
         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; \
+        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; \
-    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;
+    # Verify ensurepip can bootstrap pip. Required by boot.go worker venv creation.
+    python3 -m ensurepip;
 
 ENTRYPOINT ["/opt/apache/beam/boot"]
 
diff --git a/sdks/python/container/license_scripts/upgrade_bundled_pip.py b/sdks/python/container/license_scripts/upgrade_bundled_pip.py
index 7a2a667d398e..17a750b57029 100644
--- a/sdks/python/container/license_scripts/upgrade_bundled_pip.py
+++ b/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):