From e91649b4afc7fdb36c41667ef13bef76cc55dd6f Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Thu, 1 Aug 2024 11:52:36 -0400 Subject: [PATCH] Adds extra-index-url to default image builder Signed-off-by: Thomas J. Fan --- flytekit/image_spec/default_builder.py | 13 ++++++++++--- .../unit/core/image_spec/test_default_builder.py | 2 ++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 2b343b7d3a..4e3275c5d0 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -159,14 +159,21 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): requirements_uv_path = tmp_dir / "requirements_uv.txt" requirements_uv_path.write_text("\n".join(uv_requirements)) - pip_extra = f"--index-url {image_spec.pip_index}" if image_spec.pip_index else "" - uv_python_install_command = UV_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra) + pip_extra_args = "" + + if image_spec.pip_index: + pip_extra_args += f"--index-url {image_spec.pip_index}" + if image_spec.pip_extra_index_url: + extra_urls = [f"--extra-index-url {url}" for url in image_spec.pip_extra_index_url] + pip_extra_args += " ".join(extra_urls) + + uv_python_install_command = UV_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra_args) if pip_requirements: requirements_uv_path = tmp_dir / "requirements_pip.txt" requirements_uv_path.write_text(os.linesep.join(pip_requirements)) - pip_python_install_command = PIP_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra) + pip_python_install_command = PIP_PYTHON_INSTALL_COMMAND_TEMPLATE.substitute(PIP_EXTRA=pip_extra_args) else: pip_python_install_command = "" 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 e42337b567..6887f472b3 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -30,6 +30,7 @@ def test_create_docker_context(tmp_path): source_root=os.fspath(source_root), commands=["mkdir my_dir"], entrypoint=["/bin/bash"], + pip_extra_index_url=["https://extra-url.com"] ) create_docker_context(image_spec, docker_context_path) @@ -42,6 +43,7 @@ def test_create_docker_context(tmp_path): assert "scipy==1.13.0 numpy" in dockerfile_content assert "python=3.12" in dockerfile_content assert "--requirement requirements_uv.txt" in dockerfile_content + assert "--extra-index-url" in dockerfile_content assert "COPY --chown=flytekit ./src /root" in dockerfile_content assert "RUN mkdir my_dir" in dockerfile_content assert "ENTRYPOINT [\"/bin/bash\"]" in dockerfile_content