[Flytekit] Support extra copy commands in ImageSpec

flytekit/image_spec/default_builder.py

@@ -75,6 +75,7 @@
 
 $COPY_COMMAND_RUNTIME
 RUN $RUN_COMMANDS
+$EXTRA_COPY_CMDS
 
 WORKDIR /root
 SHELL ["/bin/bash", "-c"]
@@ -207,6 +208,14 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
     else:
         run_commands = ""
 
+    extra_copy_cmds = ""
+    if image_spec.copy_src_dest:
+        for src, dest in image_spec.copy_src_dest:
+            src_files_and_dirs = " ".join(src)
+            extra_copy_cmds += f"COPY --chown=flytekit {src_files_and_dirs} {dest}\n"
+    else:
+        extra_copy_cmds = ""
+
     docker_content = DOCKER_FILE_TEMPLATE.substitute(
         PYTHON_VERSION=python_version,
         UV_PYTHON_INSTALL_COMMAND=uv_python_install_command,
@@ -218,6 +227,7 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
         COPY_COMMAND_RUNTIME=copy_command_runtime,
         ENTRYPOINT=entrypoint,
         RUN_COMMANDS=run_commands,
+        EXTRA_COPY_CMDS=extra_copy_cmds,
     )
 
     dockerfile_path = tmp_dir / "Dockerfile"
@@ -247,6 +257,7 @@ class DefaultImageBuilder(ImageSpecBuilder):
         "pip_extra_index_url",
         # "registry_config",
         "commands",
+        "copy_src_dest",
     }
 
     def build_image(self, image_spec: ImageSpec) -> str:

flytekit/image_spec/image_spec.py

@@ -50,6 +50,7 @@ class ImageSpec:
         commands: Command to run during the building process
         tag_format: Custom string format for image tag. The ImageSpec hash passed in as `spec_hash`. For example,
             to add a "dev" suffix to the image tag, set `tag_format="{spec_hash}-dev"`
+        copy_src_dest: List of tuples containing source files/directories and their corresponding destination directory for copying.
     """
 
     name: str = "flytekit"
@@ -73,6 +74,7 @@ class ImageSpec:
     entrypoint: Optional[List[str]] = None
     commands: Optional[List[str]] = None
     tag_format: Optional[str] = None
+    copy_src_dest: Optional[List[Tuple[List[str], str]]] = None
 
     def __post_init__(self):
         self.name = self.name.lower()
@@ -232,6 +234,18 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec":
 
         return new_image_spec
 
+    def copy(self, src: List[str], dest: str = ".") -> "ImageSpec":
+        """
+        Builder that returns a new image spec with the source files copied to the destination directory.
+        """
+        new_image_spec = copy.deepcopy(self)
+        if new_image_spec.copy_src_dest is None:
+            new_image_spec.copy_src_dest = []
+
+        new_image_spec.copy_src_dest.append((src, dest))
+
+        return new_image_spec
+
     def force_push(self) -> "ImageSpec":
         """
         Builder that returns a new image spec with force push enabled.

tests/flytekit/unit/core/image_spec/test_default_builder.py

@@ -31,7 +31,8 @@ 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"]
+        pip_extra_index_url=["https://extra-url.com"],
+        copy_src_dest=[(["file1.txt", "/dir1"], "/root")]
     )
 
     create_docker_context(image_spec, docker_context_path)
@@ -49,6 +50,7 @@ def test_create_docker_context(tmp_path):
     assert "RUN mkdir my_dir" in dockerfile_content
     assert "ENTRYPOINT [\"/bin/bash\"]" in dockerfile_content
     assert "mkdir -p $HOME" in dockerfile_content
+    assert "COPY --chown=flytekit file1.txt /dir1 /root" in dockerfile_content
 
     requirements_path = docker_context_path / "requirements_uv.txt"
     assert requirements_path.exists()
@@ -177,6 +179,7 @@ def test_build(tmp_path):
         requirements=os.fspath(other_requirements_path),
         source_root=os.fspath(source_root),
         commands=["mkdir my_dir"],
+        copy_src_dest=[([f"{tmp_path}/hello_world.txt", f"{tmp_path}/requirements.txt"], "/my_dir")]
     )
 
     builder = DefaultImageBuilder()

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -26,12 +26,14 @@ def test_image_spec(mock_image_spec_builder):
         requirements=REQUIREMENT_FILE,
         registry_config=REGISTRY_CONFIG_FILE,
         entrypoint=["/bin/bash"],
+        copy_src_dest=[(["/src/file1.txt"], "/dest")]
     )
     assert image_spec._is_force_push is False
 
     image_spec = image_spec.with_commands("echo hello")
     image_spec = image_spec.with_packages("numpy")
     image_spec = image_spec.with_apt_packages("wget")
+    image_spec = image_spec.copy(["/src", "/src/file2.txt"], "/dest")
     image_spec = image_spec.force_push()
 
     assert image_spec.python_version == "3.8"
@@ -52,6 +54,7 @@ def test_image_spec(mock_image_spec_builder):
     assert image_spec.commands == ["echo hello"]
     assert image_spec._is_force_push is True
     assert image_spec.entrypoint == ["/bin/bash"]
+    assert image_spec.copy_src_dest == [(["/src/file1.txt"], "/dest"), (["/src", "/src/file2.txt"], "/dest")]
 
     tag = calculate_hash_from_image_spec(image_spec)
     assert "=" != tag[-1]