Add docker commands in ImageSpec

flytekit/image_spec/default_builder.py

@@ -77,6 +77,7 @@
 
 $COPY_COMMAND_RUNTIME
 RUN $RUN_COMMANDS
+$DOCKER_COMMANDS
 
 WORKDIR /root
 SHELL ["/bin/bash", "-c"]
@@ -210,6 +211,11 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path):
     else:
         run_commands = ""
 
+    if image_spec.docker_commands:
+        docker_commands = "\n".join(image_spec.docker_commands)
+    else:
+        docker_commands = ""
+
     docker_content = DOCKER_FILE_TEMPLATE.substitute(
         PYTHON_VERSION=python_version,
         UV_PYTHON_INSTALL_COMMAND=uv_python_install_command,
@@ -221,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,
+        DOCKER_COMMANDS=docker_commands,
     )
 
     dockerfile_path = tmp_dir / "Dockerfile"
@@ -249,6 +256,7 @@ class DefaultImageBuilder(ImageSpecBuilder):
         "pip_index",
         # "registry_config",
         "commands",
+        "docker_commands",
     }
 
     def build_image(self, image_spec: ImageSpec) -> str:

flytekit/image_spec/image_spec.py

@@ -49,6 +49,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"`
+        docker_commands: List of docker commands to run during the building process
     """
 
     name: str = "flytekit"
@@ -72,6 +73,7 @@ class ImageSpec:
     entrypoint: Optional[List[str]] = None
     commands: Optional[List[str]] = None
     tag_format: Optional[str] = None
+    docker_commands: Optional[List[str]] = None
 
     def __post_init__(self):
         self.name = self.name.lower()
@@ -87,6 +89,7 @@ def __post_init__(self):
             "pip_extra_index_url",
             "entrypoint",
             "commands",
+            "docker_commands",
         ]
         for parameter in parameters_str_list:
             attr = getattr(self, parameter)
@@ -227,6 +230,21 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec":
 
         return new_image_spec
 
+    def with_docker_commands(self, docker_commands: Union[str, List[str]]) -> "ImageSpec":
+        """
+        Builder that returns a new image spec with additional list of docker commands that will be executed during the building process.
+        """
+        new_image_spec = copy.deepcopy(self)
+        if new_image_spec.docker_commands is None:
+            new_image_spec.docker_commands = []
+
+        if isinstance(docker_commands, List):
+            new_image_spec.docker_commands.extend(docker_commands)
+        else:
+            new_image_spec.docker_commands.append(docker_commands)
+
+        return new_image_spec
+
     def force_push(self) -> "ImageSpec":
         """
         Builder that returns a new image spec with force push enabled.

plugins/flytekit-envd/flytekitplugins/envd/image_builder.py

@@ -28,6 +28,9 @@ class EnvdImageSpecBuilder(ImageSpecBuilder):
     def build_image(self, image_spec: ImageSpec):
         cfg_path = create_envd_config(image_spec)
 
+        if image_spec.docker_commands:
+            raise ValueError("Docker commands are not supported in envd builder.")
+
         if image_spec.registry_config:
             bootstrap_command = f"envd bootstrap --registry-config {image_spec.registry_config}"
             execute_command(bootstrap_command)

plugins/flytekit-envd/tests/test_image_spec.py

@@ -128,3 +128,12 @@ def build():
     )
 
     assert contents == expected_contents
+
+def test_image_spec_with_envd_builder_exception():
+    image_spec = ImageSpec(
+        name="envd_image",
+        builder="envd",
+        docker_commands=["RUN ls"],
+    )
+    with pytest.raises(ValueError):
+        ImageBuildEngine.build(image_spec)

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"],
+        docker_commands=["RUN git clone https://github.com/flyteorg/flytekit.git", "COPY . /root"],
     )
 
     create_docker_context(image_spec, docker_context_path)
@@ -49,6 +50,8 @@ 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 "RUN git clone https://github.com/flyteorg/flytekit.git" in dockerfile_content
+    assert "COPY . /root" in dockerfile_content
 
     requirements_path = docker_context_path / "requirements_uv.txt"
     assert requirements_path.exists()
@@ -171,12 +174,13 @@ def test_build(tmp_path):
         name="FLYTEKIT",
         python_version="3.12",
         env={"MY_ENV": "MY_VALUE"},
-        apt_packages=["curl"],
+        apt_packages=["curl", "git"],
         conda_packages=["scipy==1.13.0", "numpy"],
         packages=["pandas==2.2.1"],
         requirements=os.fspath(other_requirements_path),
         source_root=os.fspath(source_root),
         commands=["mkdir my_dir"],
+        docker_commands=["RUN git clone https://github.com/flyteorg/flytekit.git", "COPY . /root"],
     )
 
     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"],
+        docker_commands=["RUN ls"],
     )
     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.with_docker_commands(["RUN echo hello"])
     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.docker_commands == ["RUN ls", "RUN echo hello"]
 
     tag = calculate_hash_from_image_spec(image_spec)
     assert "=" != tag[-1]