Support uv editable installs; fix cwd editable detection and secret mounts (#560)

redartera · web-flow · commit b3f3daa274c9 · 2026-01-27T09:50:52.000-08:00
Signed-off-by: redartera &lt;reda@artera.ai&gt;
diff --git a/src/flyte/_internal/imagebuild/docker_builder.py b/src/flyte/_internal/imagebuild/docker_builder.py
@@ -12,6 +12,7 @@
 import click
 
 from flyte import Secret
+from flyte._code_bundle._ignore import STANDARD_IGNORE_PATTERNS
 from flyte._image import (
     AptPackages,
     Commands,
@@ -38,7 +39,11 @@
     LocalDockerCommandImageChecker,
     LocalPodmanCommandImageChecker,
 )
-from flyte._internal.imagebuild.utils import copy_files_to_context, get_and_list_dockerignore
+from flyte._internal.imagebuild.utils import (
+    copy_files_to_context,
+    get_and_list_dockerignore,
+    get_uv_editable_install_mounts,
+)
 from flyte._logging import logger
 
 _F_IMG_ID = "_F_IMG_ID"
@@ -49,6 +54,7 @@
 RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
    --mount=type=bind,target=uv.lock,src=$UV_LOCK_PATH,rw \
    --mount=type=bind,target=pyproject.toml,src=$PYPROJECT_PATH \
+   $EDITABLE_INSTALL_MOUNTS \
    $SECRET_MOUNT \
    VIRTUAL_ENV=$${VIRTUAL_ENV-/opt/venv} uv sync --active --inexact $PIP_INSTALL_ARGS
 """)
@@ -271,16 +277,24 @@ async def handle(
             pip_install_args = " ".join(layer.get_pip_install_args())
             if "--no-install-project" not in pip_install_args:
                 pip_install_args += " --no-install-project"
-            if "--no-sources" not in pip_install_args:
-                pip_install_args += " --no-sources"
-            # Only Copy pyproject.yaml and uv.lock.
+            # Only Copy pyproject.yaml and uv.lock from the project root.
             pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
             uvlock_dst = copy_files_to_context(layer.uvlock, context_path)
+            # Apply any editable install mounts to the template.
+            editable_install_mounts = get_uv_editable_install_mounts(
+                project_root=layer.pyproject.parent,
+                context_path=context_path,
+                ignore_patterns=[
+                    *STANDARD_IGNORE_PATTERNS,
+                    *docker_ignore_patterns,
+                ],
+            )
             delta = UV_LOCK_WITHOUT_PROJECT_INSTALL_TEMPLATE.substitute(
                 UV_LOCK_PATH=uvlock_dst.relative_to(context_path),
                 PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
                 PIP_INSTALL_ARGS=pip_install_args,
                 SECRET_MOUNT=secret_mounts,
+                EDITABLE_INSTALL_MOUNTS=editable_install_mounts,
             )
         else:
             # Copy the entire project.
@@ -429,19 +443,19 @@ def _get_secret_command(secret: str | Secret) -> typing.List[str]:
 def _get_secret_mounts_layer(secrets: typing.Tuple[str | Secret, ...] | None) -> str:
     if secrets is None:
         return ""
-    secret_mounts_layer = ""
+    secret_mounts_layer = []
     for s in secrets:
         secret = Secret(key=s) if isinstance(s, str) else s
         secret_id = hash(secret)
         if secret.mount:
-            secret_mounts_layer += f"--mount=type=secret,id={secret_id},target={secret.mount}"
+            secret_mounts_layer.append(f"--mount=type=secret,id={secret_id},target={secret.mount}")
         elif secret.as_env_var:
-            secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret.as_env_var}"
+            secret_mounts_layer.append(f"--mount=type=secret,id={secret_id},env={secret.as_env_var}")
         else:
             secret_default_env_key = "_".join(list(filter(None, (secret.group, secret.key))))
-            secret_mounts_layer += f"--mount=type=secret,id={secret_id},env={secret_default_env_key}"
+            secret_mounts_layer.append(f"--mount=type=secret,id={secret_id},env={secret_default_env_key}")
 
-    return secret_mounts_layer
+    return " ".join(secret_mounts_layer)
 
 
 async def _process_layer(
diff --git a/src/flyte/_internal/imagebuild/remote_builder.py b/src/flyte/_internal/imagebuild/remote_builder.py
@@ -13,6 +13,7 @@
 from flyteidl2.common import phase_pb2
 
 import flyte
+from flyte._code_bundle._ignore import STANDARD_IGNORE_PATTERNS
 import flyte.errors
 from flyte import Image, remote
 from flyte._code_bundle._utils import tar_strip_file_attributes
@@ -34,7 +35,11 @@
     WorkDir,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
-from flyte._internal.imagebuild.utils import copy_files_to_context, get_and_list_dockerignore
+from flyte._internal.imagebuild.utils import (
+    copy_files_to_context,
+    get_and_list_dockerignore,
+    get_uv_project_editable_dependencies,
+)
 from flyte._internal.runtime.task_serde import get_security_context
 from flyte._logging import logger
 from flyte._secret import Secret
@@ -299,25 +304,42 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             )
             layers.append(pip_layer)
         elif isinstance(layer, UVProject):
-            if layer.project_install_mode == "dependencies_only":
-                # Copy pyproject itself
-                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
-                if pip_options.extra_args:
-                    if "--no-install-project" not in pip_options.extra_args:
+            pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            # Keep track of the directory containing the pyproject.toml file
+            # this is what should be passed to the UVProject image definition proto as 'pyproject'
+            pyproject_dir_dst = pyproject_dst.parent
+
+            # Copy uv.lock itself
+            uvlock_dst = copy_files_to_context(layer.uvlock, context_path)
+
+            # Handle the project install mode
+            match layer.project_install_mode:
+                case "dependencies_only":
+                    if pip_options.extra_args and ("--no-install-project" not in pip_options.extra_args):
                         pip_options.extra_args += " --no-install-project"
-                else:
-                    pip_options.extra_args = " --no-install-project"
-                if "--no-sources" not in pip_options.extra_args:
-                    pip_options.extra_args += " --no-sources"
-            else:
-                # Copy the entire project
-                docker_ignore_patterns = get_and_list_dockerignore(image)
-                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path, docker_ignore_patterns)
+                    # Copy any editable dependencies to the context
+                    # We use the docker ignore patterns to avoid copying the editable dependencies to the context.
+                    docker_ignore_patterns = get_and_list_dockerignore(image)
+                    standard_ignore_patterns = STANDARD_IGNORE_PATTERNS.copy()
+                    for editable_dep in get_uv_project_editable_dependencies(layer.pyproject.parent):
+                        copy_files_to_context(
+                            editable_dep,
+                            context_path,
+                            ignore_patterns=[*standard_ignore_patterns, *docker_ignore_patterns],
+                        )
+                case "install_project":
+                    # Copy the entire project
+                    docker_ignore_patterns = get_and_list_dockerignore(image)
+                    pyproject_dir_dst = copy_files_to_context(layer.pyproject.parent, context_path, docker_ignore_patterns)
+                case _:
+                    raise ValueError(f"Invalid project install mode: {layer.project_install_mode}")
 
             uv_layer = image_definition_pb2.Layer(
                 uv_project=image_definition_pb2.UVProject(
-                    pyproject=str(pyproject_dst.relative_to(context_path)),
-                    uvlock=str(copy_files_to_context(layer.uvlock, context_path).relative_to(context_path)),
+                    # NOTE: UVProject expects 'pyproject' to be the directory containing the pyproject.toml file
+                    # whereas it expects 'uvlock' to be the path to the uv.lock file itself.
+                    pyproject=str(pyproject_dir_dst.relative_to(context_path)),
+                    uvlock=str(uvlock_dst.relative_to(context_path)),
                     options=pip_options,
                     secret_mounts=secret_mounts,
                 )
diff --git a/src/flyte/_internal/imagebuild/utils.py b/src/flyte/_internal/imagebuild/utils.py
@@ -1,7 +1,10 @@
+import re
 import shutil
+import subprocess
 from pathlib import Path, PurePath
 from typing import List, Optional
 
+from flyte._code_bundle._ignore import STANDARD_IGNORE_PATTERNS
 from flyte._image import DockerIgnore, Image
 from flyte._logging import logger
 
@@ -77,3 +80,65 @@ def get_and_list_dockerignore(image: Image) -> List[str]:
         logger.error(f"Failed to read .dockerignore file at {dockerignore_path}: {e}")
         return []
     return patterns
+
+
+def _extract_editables_from_uv_export(project_root: Path) -> list[str]:
+    """Extracts editable dependencies from a uv export output."""
+    uv_export = subprocess.run(
+        ["uv", "export", "--no-emit-project"], cwd=project_root, capture_output=True, text=True, check=True
+    )
+    matches = []
+    for line in uv_export.stdout.splitlines():
+        if match := re.search(r"-e\s+([^\s]+)", line):
+            matches.append(match.group(1))
+    return matches
+
+
+def get_uv_project_editable_dependencies(project_root: Path) -> list[Path]:
+    """Parses uv export output to find editable path dependencies for a given project.
+
+    Args:
+        project_root: Root of the uv project to inspect.
+
+    Returns:
+        A list of local paths referenced as editable dependencies.
+    """
+    paths = []
+    for match in _extract_editables_from_uv_export(project_root):
+        # If the the path is absolute already, keep as-is
+        # otherwise we need to complete it by pre-pending the project root where 'uv export' was run from.
+        resolved_path = Path(match) if Path(match).is_absolute() else (project_root / match)
+        # Raise an error if the path isn't a child of the project root
+        if not resolved_path.is_relative_to(project_root):
+            raise ValueError(
+                "Editable dependency paths must be within the project root, this is not supported."
+                f"Found {resolved_path=} outside of {project_root=}."
+            )
+        paths.append(resolved_path)
+    return paths
+
+
+def get_uv_editable_install_mounts(
+    project_root: Path, context_path: Path, ignore_patterns: list[str] | None = None
+) -> str:
+    """Builds Docker bind mounts for uv editable path dependencies.
+
+    Args:
+        project_root: Root of the uv project to inspect.
+        context_path: Build context directory for Docker.
+        ignore_patterns: A list of ignore patterns to apply when copying editable dependency contents.
+            If None, the standard ignore patterns of 'StandardIgnore' will be used.
+    Returns:
+        A string of Docker bind-mount arguments for editable dependencies.
+    """
+    ignore_patterns = ignore_patterns or STANDARD_IGNORE_PATTERNS.copy()
+    mounts = []
+    for editable_dep in get_uv_project_editable_dependencies(project_root):
+        # Copy the contents of the editable install by applying ignores
+        editable_dep_within_context = copy_files_to_context(editable_dep, context_path, ignore_patterns=ignore_patterns)
+        mounts.append(
+            "--mount=type=bind,"
+            f"src={editable_dep_within_context.relative_to(context_path)},"
+            f"target={editable_dep.relative_to(project_root)}"
+        )
+    return " ".join(mounts)
diff --git a/src/flyte/_utils/helpers.py b/src/flyte/_utils/helpers.py
@@ -106,16 +106,20 @@ def get_cwd_editable_install() -> typing.Optional[Path]:
     #   - if cwd is nested inside the editable install folder.
     #   - if the cwd is exactly one level above the editable install folder.
     cwd = Path.cwd()
+
+    # First, prioritize editable installs that are parents of cwd
     for install in editable_installs:
         # child.is_relative_to(parent) is True if child is inside parent
         if cwd.is_relative_to(install):
             return install
-        else:
-            # check if the cwd is one level above the install folder
-            if install.parent == cwd:
-                # check if the install folder contains a pyproject.toml file
-                if (cwd / "pyproject.toml").exists() or (cwd / "setup.py").exists():
-                    return install  # note we want the install folder, not the parent
+
+    # Then maybe see if one of the installs is exactly one level above cwd
+    for install in editable_installs:
+        # check if the cwd is one level above the install folder
+        if install.parent == cwd:
+            # check if the install folder contains a pyproject.toml file
+            if (cwd / "pyproject.toml").exists() or (cwd / "setup.py").exists():
+                return install  # note we want the install folder, not the parent
 
     return None
 
diff --git a/tests/flyte/imagebuild/test_docker_builder.py b/tests/flyte/imagebuild/test_docker_builder.py
@@ -1,11 +1,14 @@
+import asyncio
+import subprocess
 import tempfile
 from pathlib import Path, PurePath
 from unittest.mock import patch
 
 import pytest
+import pytest_asyncio
 
 from flyte import Secret
-from flyte._image import Image, PipPackages, PoetryProject, Requirements
+from flyte._image import Image, PipPackages, PoetryProject, Requirements, UVProject
 from flyte._internal.imagebuild.docker_builder import (
     CopyConfig,
     CopyConfigHandler,
@@ -429,3 +432,55 @@ async def test_uvproject_handler_with_project_install():
             assert (expected_dst_path / "uv.lock").exists(), "uv.lock should be included"
             assert not (expected_dst_path / "memo.txt").exists(), "memo.txt should be excluded"
             assert not (expected_dst_path / ".cache").exists(), ".cache directory should be excluded"
+
+
+@pytest_asyncio.fixture
+async def uv_project_with_editable(tmp_path: Path):
+    """An empty uv project with a single editable dependency"""
+
+    async def _uv(cmd: list[str], cwd: Path):
+        return await asyncio.to_thread(
+            subprocess.run, ["uv", *cmd], cwd=str(cwd), capture_output=True, text=True, check=True
+        )
+
+    project_root = tmp_path / "project"
+    project_root.mkdir(parents=True)
+    # Create a main project
+    await _uv(["init", "--lib"], project_root)
+    # Create an editable dependency
+    dep_folder = project_root / "libs" / "editable_dep"
+    dep_folder.mkdir(parents=True)
+    # Create an editable dependency project and add it to the main project
+    await _uv(["init", "--lib"], dep_folder)
+    await _uv(["add", "--editable", "./libs/editable_dep", "--no-sync"], project_root)
+    # Generate a lock file for the main project
+    await _uv(["lock"], project_root)
+    yield project_root, dep_folder
+
+
+@pytest.mark.asyncio
+async def test_uvproject_handler_includes_editable_mounts_in_dependencies_only_mode(uv_project_with_editable):
+    with tempfile.TemporaryDirectory() as tmp_context:
+        context_path = Path(tmp_context)
+
+        project_root, dep_folder = uv_project_with_editable
+        pyproject_file = project_root / "pyproject.toml"
+        uv_lock_file = project_root / "uv.lock"
+
+        uv_project = UVProject(
+            pyproject=pyproject_file.absolute(),
+            uvlock=uv_lock_file.absolute(),
+            project_install_mode="dependencies_only",
+        )
+
+        initial_dockerfile = "FROM python:3.9\n"
+        result = await UVProjectHandler.handle(
+            layer=uv_project,
+            context_path=context_path,
+            dockerfile=initial_dockerfile,
+            docker_ignore_patterns=[],
+        )
+        expected_dep_in_context = "_flyte_abs_context" + str(dep_folder)
+        expected_dep_in_container = dep_folder.relative_to(project_root)
+        expected_mount = f"--mount=type=bind,src={expected_dep_in_context},target={expected_dep_in_container}"
+        assert expected_mount in result
diff --git a/tests/flyte/imagebuild/test_utils.py b/tests/flyte/imagebuild/test_utils.py
