From cacb9f0eee9beb7997c57e0faded0419e1afb53f Mon Sep 17 00:00:00 2001 From: Jason Lai Date: Wed, 6 Mar 2024 16:56:26 +0800 Subject: [PATCH] Remove trailing '=' character from generated image tag hash value (#2240) Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 6 +++--- tests/flytekit/unit/cli/pyflyte/test_run.py | 4 ++-- tests/flytekit/unit/core/image_spec/test_image_spec.py | 1 + 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 7a8ef547da..54f0618c2d 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -257,9 +257,9 @@ def calculate_hash_from_image_spec(image_spec: ImageSpec): # won't rebuild the image if we change the registry_config path spec.registry_config = None image_spec_bytes = asdict(spec).__str__().encode("utf-8") - tag = base64.urlsafe_b64encode(hashlib.md5(image_spec_bytes).digest()).decode("ascii") - # replace "=" with "." and replace "-" with "_" to make it a valid tag - return tag.replace("=", ".").replace("-", "_") + tag = base64.urlsafe_b64encode(hashlib.md5(image_spec_bytes).digest()).decode("ascii").rstrip("=") + # replace "-" with "_" to make it a valid tag + return tag.replace("-", "_") def hash_directory(path): diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index b42be184f8..9dac0c8e8b 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -297,9 +297,9 @@ def test_list_default_arguments(wf_path): ) ic_result_4 = ImageConfig( - default_image=Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."), + default_image=Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q"), images=[ - Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."), + Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q"), Image(name="xyz", fqn="docker.io/xyz", tag="latest"), Image(name="abc", fqn="docker.io/abc", tag=None), ], 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 38727d02fd..9ec1adc77c 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -49,6 +49,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.commands == ["echo hello"] tag = calculate_hash_from_image_spec(image_spec) + assert "=" != tag[-1] assert image_spec.image_name() == f"flytekit:{tag}" ctx = context_manager.FlyteContext.current_context() with context_manager.FlyteContextManager.with_context(