Adds extra-index-url to default image builder (#2636)

Signed-off-by: Thomas J. Fan <[email protected]>
Co-authored-by: Kevin Su <[email protected]>

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 = ""
 

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