From 699749695e6e99702efffe13817c0f07c8451713 Mon Sep 17 00:00:00 2001 From: Hongxin Liang Date: Tue, 16 Jun 2026 15:40:10 +0200 Subject: [PATCH] Fix template replacement producing invalid image reference with digests When an image with a digest is used (e.g. image:tag@sha256:...), the template replacement for {{.image..version}} hardcodes : as the separator. Since version returns the digest, this produces image:tag:sha256:... (invalid) instead of image:tag@sha256:... (valid). Fix by detecting when version is a digest (contains :) and replacing the preceding : separator with @ in the output. Fixes https://github.com/flyteorg/flyte/issues/7531 Signed-off-by: Hongxin Liang --- flytekit/core/python_auto_container.py | 12 +++++++++--- .../unit/core/test_python_function_task.py | 18 +++++++++++++++++- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 0e3dd5bda8..ace756ca0b 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -467,10 +467,16 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: if img_cfg is None: raise AssertionError(f"Image Config with name {name} not found in the configuration") if attr == "version": - if img_cfg.version is not None: - img = img.replace(replace_group, img_cfg.version) + version = img_cfg.version if img_cfg.version is not None else cfg.default_image.version + if ":" in version: + # version is a digest (e.g. sha256:...) — needs @ separator, not : + colon_prefixed = f":{replace_group}" + if colon_prefixed in img: + img = img.replace(colon_prefixed, f"@{version}") + else: + img = img.replace(replace_group, version) else: - img = img.replace(replace_group, cfg.default_image.version) + img = img.replace(replace_group, version) elif attr == "fqn": img = img.replace(replace_group, img_cfg.fqn) elif attr == "": diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 427e3bd774..b2e4d9c1b9 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -58,7 +58,13 @@ def test_container_image_conversion(mock_image_spec_builder): fqn="xyz.com/other5", tag="tag5" ) - cfg = ImageConfig(default_image=default_img, images=[default_img, other_img, other_img2, other_img3, other_img4, other_img5]) + # Simulates _parse_image_identifier("xyz.com/other6:tag6@sha256:...") where the tag is baked into fqn + other_img6 = Image( + name="other6", + fqn="xyz.com/other6:tag6", + digest="sha256:26c68657ccce2cb0a31b330cb0be2b5e108d467f641c62e13ab40cbec258c68d", + ) + cfg = ImageConfig(default_image=default_img, images=[default_img, other_img, other_img2, other_img3, other_img4, other_img5, other_img6]) assert get_registerable_container_image(None, cfg) == "xyz.com/abc:tag1" assert get_registerable_container_image("", cfg) == "xyz.com/abc:tag1" assert get_registerable_container_image("abc", cfg) == "abc" @@ -78,6 +84,16 @@ def test_container_image_conversion(mock_image_spec_builder): get_registerable_container_image("{{.image.other2.fqn}}@{{.image.other2.version}}", cfg) == "xyz.com/other2@sha256:26c68657ccce2cb0a31b330cb0be2b5e108d467f641c62e13ab40cbec258c68d" ) + # Digest image with : separator in template should be fixed to @ + assert ( + get_registerable_container_image("{{.image.other2.fqn}}:{{.image.other2.version}}", cfg) + == "xyz.com/other2@sha256:26c68657ccce2cb0a31b330cb0be2b5e108d467f641c62e13ab40cbec258c68d" + ) + # repo:tag@sha256:digest — fqn already contains the tag, digest needs @ separator + assert ( + get_registerable_container_image("{{.image.other6.fqn}}:{{.image.other6.version}}", cfg) + == "xyz.com/other6:tag6@sha256:26c68657ccce2cb0a31b330cb0be2b5e108d467f641c62e13ab40cbec258c68d" + ) assert ( get_registerable_container_image("{{.image.other3.fqn}}:{{.image.other3.version}}", cfg) == "xyz.com/other3:tag1"