diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 027b5e0b47..b4772d1a92 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -46,6 +46,8 @@ class ImageSpec: 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 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"` """ name: str = "flytekit" @@ -67,6 +69,7 @@ class ImageSpec: pip_extra_index_url: Optional[List[str]] = None registry_config: Optional[str] = None commands: Optional[List[str]] = None + tag_format: Optional[str] = None def __post_init__(self): self.name = self.name.lower() @@ -85,6 +88,9 @@ def image_name(self) -> str: def _image_name(self) -> str: """Construct full image name with tag.""" tag = calculate_hash_from_image_spec(self) + if self.tag_format: + tag = self.tag_format.format(spec_hash=tag) + container_image = f"{self.name}:{tag}" if self.registry: container_image = f"{self.registry}/{container_image}" 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 f1b0bbe975..15b14f6687 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -112,3 +112,13 @@ def test_build_existing_image_with_force_push(): ImageBuildEngine.build(image_spec) ImageBuildEngine._build_image.assert_called_once() + + +def test_custom_tag(): + spec = ImageSpec( + name="my_image", + python_version="3.11", + tag_format="{spec_hash}-dev", + ) + spec_hash = calculate_hash_from_image_spec(spec) + assert spec.image_name() == f"my_image:{spec_hash}-dev"