diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 9e37e9c8c5..96b9acbe9e 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -35,6 +35,31 @@ """ ) +POETRY_LOCK_TEMPLATE = Template( + """\ +RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \ + --mount=from=uv,source=/uv,target=/usr/bin/uv \ + uv pip install poetry + +ENV POETRY_CACHE_DIR=/tmp/poetry_cache \ + POETRY_VIRTUALENVS_IN_PROJECT=true + +# poetry install does not work running in /, so we move to /root to create the venv +WORKDIR /root + +RUN --mount=type=cache,sharing=locked,mode=0777,target=/tmp/poetry_cache,id=poetry \ + --mount=type=bind,target=poetry.lock,src=poetry.lock \ + --mount=type=bind,target=pyproject.toml,src=pyproject.toml \ + poetry install $PIP_INSTALL_ARGS + +WORKDIR / + +# Update PATH and UV_PYTHON to point to venv +ENV PATH="/root/.venv/bin:$$PATH" \ + UV_PYTHON=/root/.venv/bin/python +""" +) + UV_PYTHON_INSTALL_COMMAND_TEMPLATE = Template( """\ RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \ @@ -44,6 +69,7 @@ """ ) + APT_INSTALL_COMMAND_TEMPLATE = Template("""\ RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \ apt-get update && apt-get install -y --no-install-recommends \ @@ -128,29 +154,33 @@ def _is_flytekit(package: str) -> bool: return name == "flytekit" -def prepare_uv_lock_command(image_spec: ImageSpec, pip_install_args: List[str], tmp_dir: Path) -> str: - # uv sync is experimental, so our uv.lock support is also experimental - # the parameters we pass into install args could be different - warnings.warn("uv.lock support is experimental", UserWarning) - +def _copy_lock_files_into_context(image_spec: ImageSpec, lock_file: str, tmp_dir: Path): if image_spec.packages is not None: - msg = "Support for uv.lock files and packages is mutually exclusive" + msg = f"Support for {lock_file} files and packages is mutually exclusive" raise ValueError(msg) - uv_lock_path = tmp_dir / "uv.lock" - shutil.copy2(image_spec.requirements, uv_lock_path) + lock_path = tmp_dir / lock_file + shutil.copy2(image_spec.requirements, lock_path) - # uv.lock requires pyproject.toml to be included + # lock requires pyproject.toml to be included pyproject_toml_path = tmp_dir / "pyproject.toml" dir_name = os.path.dirname(image_spec.requirements) pyproject_toml_src = os.path.join(dir_name, "pyproject.toml") if not os.path.exists(pyproject_toml_src): - msg = "To use uv.lock, a pyproject.toml must be in the same directory as the lock file" + msg = f"To use {lock_file}, a pyproject.toml file must be in the same directory as the lock file" raise ValueError(msg) shutil.copy2(pyproject_toml_src, pyproject_toml_path) + +def prepare_uv_lock_command(image_spec: ImageSpec, pip_install_args: List[str], tmp_dir: Path) -> str: + # uv sync is experimental, so our uv.lock support is also experimental + # the parameters we pass into install args could be different + warnings.warn("uv.lock support is experimental", UserWarning) + + _copy_lock_files_into_context(image_spec, "uv.lock", tmp_dir) + # --locked: Assert that the `uv.lock` will remain unchanged # --no-dev: Omit the development dependency group # --no-install-project: Do not install the current project @@ -160,6 +190,15 @@ def prepare_uv_lock_command(image_spec: ImageSpec, pip_install_args: List[str], return UV_LOCK_INSTALL_TEMPLATE.substitute(PIP_INSTALL_ARGS=pip_install_args) +def prepare_poetry_lock_command(image_spec: ImageSpec, pip_install_args: List[str], tmp_dir: Path) -> str: + _copy_lock_files_into_context(image_spec, "poetry.lock", tmp_dir) + + # --no-root: Do not install the current project + pip_install_args.extend(["--no-root"]) + pip_install_args = " ".join(pip_install_args) + return POETRY_LOCK_TEMPLATE.substitute(PIP_INSTALL_ARGS=pip_install_args) + + def prepare_python_install(image_spec: ImageSpec, tmp_dir: Path) -> str: pip_install_args = [] if image_spec.pip_index: @@ -174,6 +213,8 @@ def prepare_python_install(image_spec: ImageSpec, tmp_dir: Path) -> str: requirement_basename = os.path.basename(image_spec.requirements) if requirement_basename == "uv.lock": return prepare_uv_lock_command(image_spec, pip_install_args, tmp_dir) + elif requirement_basename == "poetry.lock": + return prepare_poetry_lock_command(image_spec, pip_install_args, tmp_dir) # Assume this is a requirements.txt file with open(image_spec.requirements) as f: 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 b7d8f3af00..0abd3d9467 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -251,44 +251,70 @@ def test_create_docker_context_uv_lock(tmp_path): ) in dockerfile_content +@pytest.mark.parametrize("lock_file", ["uv.lock", "poetry.lock"]) @pytest.mark.filterwarnings("ignore::UserWarning") -def test_uv_lock_errors_no_pyproject_toml(monkeypatch, tmp_path): +def test_lock_errors_no_pyproject_toml(monkeypatch, tmp_path, lock_file): run_mock = Mock() monkeypatch.setattr("flytekit.image_spec.default_builder.run", run_mock) - uv_lock_file = tmp_path / "uv.lock" - uv_lock_file.write_text("this is a lock file") + lock_file_path = tmp_path / lock_file + lock_file_path.write_text("this is a lock file") image_spec = ImageSpec( name="FLYTEKIT", python_version="3.12", - requirements=os.fspath(uv_lock_file), + requirements=os.fspath(lock_file_path), ) builder = DefaultImageBuilder() - with pytest.raises(ValueError, match="To use uv.lock"): + with pytest.raises(ValueError, match="a pyproject.toml file must be in the same"): builder.build_image(image_spec) +@pytest.mark.parametrize("lock_file", ["uv.lock", "poetry.lock"]) @pytest.mark.filterwarnings("ignore::UserWarning") -@pytest.mark.parametrize("invalid_param", ["packages"]) -def test_uv_lock_error_no_packages(monkeypatch, tmp_path, invalid_param): +def test_uv_lock_error_no_packages(monkeypatch, tmp_path, lock_file): run_mock = Mock() monkeypatch.setattr("flytekit.image_spec.default_builder.run", run_mock) - uv_lock_file = tmp_path / "uv.lock" - uv_lock_file.write_text("this is a lock file") + lock_file_path = tmp_path / lock_file + lock_file_path.write_text("this is a lock file") image_spec = ImageSpec( name="FLYTEKIT", python_version="3.12", - requirements=os.fspath(uv_lock_file), + requirements=os.fspath(lock_file), packages=["ruff"], ) builder = DefaultImageBuilder() - with pytest.raises(ValueError, match="Support for uv.lock files and packages is mutually exclusive"): + with pytest.raises(ValueError, match=f"Support for {lock_file} files and packages is mutually exclusive"): builder.build_image(image_spec) run_mock.assert_not_called() + + +def test_create_poetry_lock(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + poetry_lock = tmp_path / "poetry.lock" + poetry_lock.write_text("this is a lock file") + + pyproject_file = tmp_path / "pyproject.toml" + pyproject_file.write_text("this is a pyproject.toml file") + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + requirements=os.fspath(poetry_lock), + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + dockerfile_content = dockerfile_path.read_text() + + assert "poetry install --no-root" in dockerfile_content