diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 3b35214c22..b4be8fda97 100644 --- a/flytekit/image_spec/default_builder.py +++ b/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: diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index e750cc211e..2ca7fe1984 100644 --- a/flytekit/image_spec/image_spec.py +++ b/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. diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 7a9f3ad955..14b41f167c 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/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) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index cbd1eb761d..fbdc321799 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/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) 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 e61a3cb7c8..f529faf8d0 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/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() diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index fa63f08993..be3233190e 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/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]