From 2c185bf8b8fb90e3f78d7b9275922f803175de5a Mon Sep 17 00:00:00 2001
From: Niels Bantilan <niels.bantilan@gmail.com>
Date: Wed, 16 Jul 2025 15:40:17 -0400
Subject: [PATCH] docker build: make sure venvs are owned by flytekit

Signed-off-by: Niels Bantilan <niels.bantilan@gmail.com>
---
 flytekit/image_spec/default_builder.py | 27 ++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py
index ac8256e862..b925ddd15b 100644
--- a/flytekit/image_spec/default_builder.py
+++ b/flytekit/image_spec/default_builder.py
@@ -29,7 +29,8 @@
     --mount=type=bind,target=uv.lock,src=uv.lock \
     --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
     $PIP_SECRET_MOUNT \
-    uv sync $PIP_INSTALL_ARGS
+    uv sync $PIP_INSTALL_ARGS && \
+    chown -R flytekit /root/.venv
 WORKDIR /
 
 # Update PATH and UV_PYTHON to point to the venv created by uv sync
@@ -54,12 +55,12 @@
     --mount=type=bind,target=poetry.lock,src=poetry.lock \
     --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
     $PIP_SECRET_MOUNT \
-    poetry install $PIP_INSTALL_ARGS
-
+    poetry install $PIP_INSTALL_ARGS && \
+    chown -R flytekit /root/.venv
 WORKDIR /
 
 # Update PATH and UV_PYTHON to point to venv
-ENV PATH="/root/.venv/bin:$$PATH" \
+ENV PATH="/root/.venv/bin:$$PATH"  \
     UV_PYTHON=/root/.venv/bin/python
 """
 )
@@ -81,6 +82,7 @@
     $APT_PACKAGES
 """)
 
+# make sure that micromamba python installation is owned by flytekit user
 MICROMAMBA_INSTALL_COMMAND_TEMPLATE = Template("""\
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/opt/micromamba/pkgs,\
 id=micromamba \
@@ -91,7 +93,8 @@
     python=$PYTHON_VERSION $CONDA_PACKAGES \
     || micromamba install -n runtime --root-prefix /opt/micromamba \
     -c conda-forge $CONDA_CHANNELS \
-    python=$PYTHON_VERSION $CONDA_PACKAGES )
+    python=$PYTHON_VERSION $CONDA_PACKAGES ) && \
+    chown -R flytekit /opt/micromamba
 """)
 
 DOCKER_FILE_TEMPLATE = Template("""\
@@ -108,6 +111,9 @@
     [ -f /etc/ssl/certs/ca-certificates.crt ] || \
     mkdir -p /etc/ssl/certs/ && cp /tmp/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
 
+RUN id -u flytekit || useradd --create-home --shell /bin/bash flytekit
+RUN chown -R flytekit /root && chown -R flytekit /home
+
 $INSTALL_PYTHON_TEMPLATE
 
 # Configure user space
@@ -119,7 +125,7 @@
     SSL_CERT_DIR=/etc/ssl/certs \
     $ENV
 
-$UV_PYTHON_INSTALL_COMMAND
+$PYTHON_INSTALL_COMMAND
 
 # Adds nvidia just in case it exists
 ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \
@@ -134,9 +140,6 @@
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
     --mount=from=uv,source=/uv,target=/usr/bin/uv $RUN_COMMANDS
 
-RUN id -u flytekit || useradd --create-home --shell /bin/bash flytekit
-RUN chown -R flytekit /root && chown -R flytekit /home
-
 WORKDIR /root
 SHELL ["/bin/bash", "-c"]
 
@@ -340,7 +343,7 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
         )
         raise ValueError(msg)
 
-    uv_python_install_command = prepare_python_install(image_spec, tmp_dir)
+    python_install_command = prepare_python_install(image_spec, tmp_dir)
     env_dict = {"PYTHONPATH": "/root"}
 
     if image_spec.env:
@@ -426,11 +429,11 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
     _f_img_id_env = f"{_F_IMG_ID}={image_spec.id}"
 
     docker_content = DOCKER_FILE_TEMPLATE.substitute(
-        UV_PYTHON_INSTALL_COMMAND=uv_python_install_command,
-        APT_INSTALL_COMMAND=apt_install_command,
         INSTALL_PYTHON_TEMPLATE=python_install_template.template,
         EXTRA_PATH=python_install_template.extra_path,
         PYTHON_EXEC=python_install_template.python_exec,
+        APT_INSTALL_COMMAND=apt_install_command,
+        PYTHON_INSTALL_COMMAND=python_install_command,
         BASE_IMAGE=base_image,
         ENV=env,
         _F_IMG_ID_ENV=_f_img_id_env,