diff --git a/flytekit/image_spec/default_builder.py b/flytekit/image_spec/default_builder.py index 7e91eebf10..a09f445ba9 100644 --- a/flytekit/image_spec/default_builder.py +++ b/flytekit/image_spec/default_builder.py @@ -1,3 +1,4 @@ +import json import os import shutil import subprocess @@ -73,6 +74,8 @@ ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \ LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$$LD_LIBRARY_PATH" +$ENTRYPOINT + $COPY_COMMAND_RUNTIME RUN $RUN_COMMANDS @@ -191,6 +194,11 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): else: python_version = f"{sys.version_info.major}.{sys.version_info.minor}" + if image_spec.entrypoint is None: + entrypoint = "" + else: + entrypoint = f"ENTRYPOINT {json.dumps(image_spec.entrypoint)}" + if image_spec.commands: run_commands = " && ".join(image_spec.commands) else: @@ -206,6 +214,7 @@ def create_docker_context(image_spec: ImageSpec, tmp_dir: Path): BASE_IMAGE=base_image, ENV=env, COPY_COMMAND_RUNTIME=copy_command_runtime, + ENTRYPOINT=entrypoint, RUN_COMMANDS=run_commands, ) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index a17a2e272e..7cde9fa70a 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -45,6 +45,7 @@ class ImageSpec: pip_index: Specify the custom pip index url pip_extra_index_url: Specify one or more pip index urls as a list registry_config: Specify the path to a JSON registry config file + entrypoint: List of strings to overwrite the entrypoint of the base image with, set to [] to remove the entrypoint. 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"` @@ -68,6 +69,7 @@ class ImageSpec: pip_index: Optional[str] = None pip_extra_index_url: Optional[List[str]] = None registry_config: Optional[str] = None + entrypoint: Optional[List[str]] = None commands: Optional[List[str]] = None tag_format: Optional[str] = None 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 bd47b4558e..8941da81b9 100644 --- a/tests/flytekit/unit/core/image_spec/test_default_builder.py +++ b/tests/flytekit/unit/core/image_spec/test_default_builder.py @@ -28,6 +28,7 @@ def test_create_docker_context(tmp_path): requirements=os.fspath(other_requirements_path), source_root=os.fspath(source_root), commands=["mkdir my_dir"], + entrypoint=["/bin/bash"], ) create_docker_context(image_spec, docker_context_path) @@ -42,6 +43,7 @@ def test_create_docker_context(tmp_path): assert "--requirement requirements_uv.txt" 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 requirements_path = docker_context_path / "requirements_uv.txt" assert requirements_path.exists() @@ -79,6 +81,24 @@ def test_create_docker_context_with_git_subfolder(tmp_path): assert requirements_path.exists() +def test_create_docker_context_with_null_entrypoint(tmp_path): + docker_context_path = tmp_path / "builder_root" + docker_context_path.mkdir() + + image_spec = ImageSpec( + name="FLYTEKIT", + python_version="3.12", + entrypoint=[], + ) + + create_docker_context(image_spec, docker_context_path) + + dockerfile_path = docker_context_path / "Dockerfile" + assert dockerfile_path.exists() + dockerfile_content = dockerfile_path.read_text() + assert "ENTRYPOINT []" in dockerfile_content + + def test_create_docker_context_cuda(tmp_path): docker_context_path = tmp_path / "builder_root" docker_context_path.mkdir() 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 02c734b119..c1e52953bb 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -25,6 +25,7 @@ def test_image_spec(mock_image_spec_builder): cudnn="8", requirements=REQUIREMENT_FILE, registry_config=REGISTRY_CONFIG_FILE, + entrypoint=["/bin/bash"], ) assert image_spec._is_force_push is False @@ -50,6 +51,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.is_container() is True assert image_spec.commands == ["echo hello"] assert image_spec._is_force_push is True + assert image_spec.entrypoint == ["/bin/bash"] tag = calculate_hash_from_image_spec(image_spec) assert "=" != tag[-1]