Added entrypoint to imagespec and default builder (#2553)

* Added entrypoint to imagespec and default builder

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

* Added docstring

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

* Revert reqs.txt

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

* Fixed falsey behavior of empty list, moved ENTRYPOINT directive into optional str to avoid entrypoint reset behavior

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

---------

Signed-off-by: pryce-turner <pryce.turner@gmail.com>

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,
     )
 

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
 

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()

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]