Add docker commands in ImageSpec

flytekit/image_spec/default_builder.py

@@ -80,6 +80,7 @@
 
 $COPY_COMMAND_RUNTIME
 RUN $RUN_COMMANDS
+$DOCKER_COMMANDS
 
 WORKDIR /root
 SHELL ["/bin/bash", "-c"]
@@ -232,6 +233,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,
@@ -244,6 +250,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"
@@ -272,6 +279,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()
@@ -227,6 +229,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.