diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index b925ddd15b..3284191076 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -75,6 +75,23 @@ """ ) +# When python_exec points to system Python (e.g. on NVIDIA base images), PEP 668 prevents +# installing packages directly. We create a venv using the base Python and install into it. +# --system-site-packages lets the venv see base image packages (e.g. NVIDIA TensorRT). +UV_PYTHON_VENV_INSTALL_COMMAND_TEMPLATE = Template( + """\ +WORKDIR /root +RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \ + --mount=from=uv,source=/uv,target=/usr/bin/uv \ + --mount=type=bind,target=requirements_uv.txt,src=requirements_uv.txt \ + $PIP_SECRET_MOUNT \ + uv venv --python $BASE_PYTHON_EXEC --system-site-packages /root/.venv && \ + uv pip install --python /root/.venv/bin/python $PIP_INSTALL_ARGS && \ + chown -R flytekit /root/.venv +WORKDIR / +""" +) + APT_INSTALL_COMMAND_TEMPLATE = Template("""\ RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \ @@ -262,7 +279,14 @@ def prepare_python_install(image_spec: ImageSpec, tmp_dir: Path) -> str: requirements.extend([line.strip() for line in f.readlines()]) if template is None: - template = UV_PYTHON_INSTALL_COMMAND_TEMPLATE + if image_spec.python_exec: + # Use venv when python_exec is set: base images with PEP 668 (externally managed + # Python) like NVIDIA TensorRT prevent direct system installs. Creating a venv + # preserves base image packages while allowing flytekit deps to be installed. + template = UV_PYTHON_VENV_INSTALL_COMMAND_TEMPLATE + else: + template = UV_PYTHON_INSTALL_COMMAND_TEMPLATE + if image_spec.packages: requirements.extend(image_spec.packages) @@ -276,10 +300,14 @@ def prepare_python_install(image_spec: ImageSpec, tmp_dir: Path) -> str: pip_install_args = " ".join(pip_install_args) - return template.substitute( + substitute_kwargs = dict( PIP_INSTALL_ARGS=pip_install_args, PIP_SECRET_MOUNT=pip_secret_mount, ) + if image_spec.python_exec and template == UV_PYTHON_VENV_INSTALL_COMMAND_TEMPLATE: + substitute_kwargs["BASE_PYTHON_EXEC"] = image_spec.python_exec + + return template.substitute(**substitute_kwargs) class _PythonInstallTemplate(NamedTuple): @@ -294,7 +322,13 @@ def prepare_python_executable(image_spec: ImageSpec) -> _PythonInstallTemplate: raise ValueError("conda_channels is not supported with python_exec") if image_spec.conda_packages: raise ValueError("conda_packages is not supported with python_exec") - return _PythonInstallTemplate(python_exec=image_spec.python_exec, template="", extra_path="") + # Packages are installed into /root/.venv (see UV_PYTHON_VENV_INSTALL_COMMAND_TEMPLATE) + # so runtime must use the venv interpreter + return _PythonInstallTemplate( + python_exec="/root/.venv/bin/python", + template="", + extra_path="/root/.venv/bin", + ) conda_packages = image_spec.conda_packages or [] conda_channels = image_spec.conda_channels or [] @@ -478,6 +512,7 @@ class DefaultImageBuilder(ImageSpecBuilder): # "registry_config", "commands", "copy", + "python_exec", "builder_config", } diff --git a/tests/flytekit/unit/core/image_spec/test_default_builder.py b/tests/flytekit/unit/core/image_spec/test_default_builder.py index 42a55220fb..366fe6e944 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -338,6 +338,7 @@ def test_create_poetry_lock(tmp_path): def test_python_exec(tmp_path): + """python_exec uses venv to avoid PEP 668 (externally managed Python) errors on NVIDIA etc.""" docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() base_image = "ghcr.io/flyteorg/flytekit:py3.11-1.14.4" @@ -355,7 +356,11 @@ def test_python_exec(tmp_path): assert dockerfile_path.exists() dockerfile_content = dockerfile_path.read_text() - assert f"UV_PYTHON={python_exec}" in dockerfile_content + # Packages install into venv; runtime uses venv Python (fixes PEP 668 on NVIDIA etc.) + assert "uv venv --python /usr/local/bin/python --system-site-packages /root/.venv" in dockerfile_content + assert "uv pip install --python /root/.venv/bin/python" in dockerfile_content + assert "UV_PYTHON=/root/.venv/bin/python" in dockerfile_content + assert "/root/.venv/bin" in dockerfile_content @pytest.mark.parametrize("key, value", [("conda_packages", ["ruff"]), ("conda_channels", ["bioconda"])])