From 5ba105dd6db7bf8e0cef5a79231e2b903f011b20 Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Sun, 3 Mar 2024 17:55:49 +0800 Subject: [PATCH 01/12] feat: implement force push functionality in ImageSpec and ImageBuildEngine classes - Add a new attribute `is_force_push` to the `ImageSpec` class - Add a new method `force_push` to the `ImageSpec` class - Update the logic in the `ImageBuildEngine` class to handle `is_force_push` attribute in `ImageSpec` instances Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 7a8ef547da..7828ace342 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -44,6 +44,7 @@ class ImageSpec: pip_index: Specify the custom pip index url registry_config: Specify the path to a JSON registry config file commands: Command to run during the building process + is_force_push: forcefully pushed to the registry """ name: str = "flytekit" @@ -64,6 +65,7 @@ class ImageSpec: pip_index: Optional[str] = None registry_config: Optional[str] = None commands: Optional[List[str]] = None + is_force_push: Optional[bool] = None def __post_init__(self): self.name = self.name.lower() @@ -181,6 +183,15 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec": return new_image_spec + def force_push(self, is_force_push: Optional[bool] = True) -> "ImageSpec": + """ + Builder that returns a new image spec with force push enabled. + """ + new_image_spec = copy.deepcopy(self) + new_image_spec.is_force_push = is_force_push + + return new_image_spec + class ImageSpecBuilder: @abstractmethod @@ -222,10 +233,13 @@ def build(cls, image_spec: ImageSpec) -> str: builder = image_spec.builder img_name = image_spec.image_name() - if img_name in cls._BUILT_IMAGES or image_spec.exist(): + if (img_name in cls._BUILT_IMAGES or image_spec.exist()) and not image_spec.is_force_push: click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: - click.secho(f"Image {img_name} not found. Building...", fg="blue") + if image_spec.is_force_push: + click.secho(f"Image {img_name} found. Force pushed image already exists.", fg="blue") + else: + click.secho(f"Image {img_name} not found. Building...", fg="blue") if builder not in cls._REGISTRY: raise Exception(f"Builder {builder} is not registered.") if builder == "envd": From 8a0ec90bd104adae5f9eb7f8b2500975f45fcb43 Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Sun, 3 Mar 2024 20:01:00 +0800 Subject: [PATCH 02/12] - feat: implement force push functionality in image_spec class - Add a `force_push` method to the `image_spec` class - Assert that `is_force_push` is True in the test function Signed-off-by: jason.lai --- tests/flytekit/unit/core/image_spec/test_image_spec.py | 2 ++ 1 file changed, 2 insertions(+) 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..be27330ee3 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -30,6 +30,7 @@ def test_image_spec(mock_image_spec_builder): image_spec = image_spec.with_commands("echo hello") image_spec = image_spec.with_packages("numpy") image_spec = image_spec.with_apt_packages("wget") + image_spec = image_spec.force_push() assert image_spec.python_version == "3.8" assert image_spec.base_image == "cr.flyte.org/flyteorg/flytekit:py3.8-latest" @@ -47,6 +48,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.pip_index is None assert image_spec.is_container() is True assert image_spec.commands == ["echo hello"] + assert image_spec.is_force_push is True tag = calculate_hash_from_image_spec(image_spec) assert image_spec.image_name() == f"flytekit:{tag}" From 4c2c39647330f826c3155f65ffbfaa1d4fd36b4f Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Mon, 4 Mar 2024 10:09:14 +0800 Subject: [PATCH 03/12] refactor: improve handling of force push in image builds - Add a default value for `is_force_push` in the `ImageSpec` class - Change the message when force pushing an image in the `ImageBuildEngine` class Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 7828ace342..dad8694fb3 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -69,6 +69,7 @@ class ImageSpec: def __post_init__(self): self.name = self.name.lower() + self.is_force_push = False # False by default if self.registry: self.registry = self.registry.lower() @@ -237,7 +238,7 @@ def build(cls, image_spec: ImageSpec) -> str: click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: if image_spec.is_force_push: - click.secho(f"Image {img_name} found. Force pushed image already exists.", fg="blue") + click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") else: click.secho(f"Image {img_name} not found. Building...", fg="blue") if builder not in cls._REGISTRY: From aa364e319f044bd63e21093092d9a97c015a526c Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Mon, 4 Mar 2024 18:19:46 +0800 Subject: [PATCH 04/12] chore: update default image tag variable in ic_result_4 - Update the default image tag in `ic_result_4` variable Signed-off-by: jason.lai --- tests/flytekit/unit/cli/pyflyte/test_run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index b42be184f8..0b2a286f94 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="guqLm2fbPHgwk0EiXmXI9g.."), images=[ - Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."), + Image(name="default", fqn="flytekit", tag="guqLm2fbPHgwk0EiXmXI9g.."), Image(name="xyz", fqn="docker.io/xyz", tag="latest"), Image(name="abc", fqn="docker.io/abc", tag=None), ], From 057ff60224b6be1731d8991466f2d63f465defbb Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Mon, 4 Mar 2024 18:36:53 +0800 Subject: [PATCH 05/12] refactor: refactor `is_force_push` attribute in ImageSpec and ImageBuildEngine classes - Change `is_force_push` attribute to `_is_force_push` in the `ImageSpec` class - Remove the `is_force_push` parameter from the `force_push` method in the `ImageSpec` class - Update references to `is_force_push` to `_is_force_push` in the `ImageBuildEngine` class Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index dad8694fb3..c5f93bb81b 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -65,11 +65,10 @@ class ImageSpec: pip_index: Optional[str] = None registry_config: Optional[str] = None commands: Optional[List[str]] = None - is_force_push: Optional[bool] = None def __post_init__(self): self.name = self.name.lower() - self.is_force_push = False # False by default + self._is_force_push = False # False by default if self.registry: self.registry = self.registry.lower() @@ -184,12 +183,12 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec": return new_image_spec - def force_push(self, is_force_push: Optional[bool] = True) -> "ImageSpec": + def force_push(self) -> "ImageSpec": """ Builder that returns a new image spec with force push enabled. """ new_image_spec = copy.deepcopy(self) - new_image_spec.is_force_push = is_force_push + self._is_force_push = True return new_image_spec @@ -234,10 +233,10 @@ def build(cls, image_spec: ImageSpec) -> str: builder = image_spec.builder img_name = image_spec.image_name() - if (img_name in cls._BUILT_IMAGES or image_spec.exist()) and not image_spec.is_force_push: + if (img_name in cls._BUILT_IMAGES or image_spec.exist()) and not image_spec._is_force_push: click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: - if image_spec.is_force_push: + if image_spec._is_force_push: click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") else: click.secho(f"Image {img_name} not found. Building...", fg="blue") From a688fbe18677cf652b1e8941586f64c82faddfeb Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Mon, 4 Mar 2024 19:00:23 +0800 Subject: [PATCH 06/12] feat: enable force push in ImageSpec class and update tests - Add a method to enable force push in `ImageSpec` class - Update image tags in test cases in `test_run.py` - Assert that `_is_force_push` is False in `test_image_spec.py` Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 1 + tests/flytekit/unit/cli/pyflyte/test_run.py | 4 ++-- tests/flytekit/unit/core/image_spec/test_image_spec.py | 3 ++- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index c5f93bb81b..28a3062222 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -188,6 +188,7 @@ def force_push(self) -> "ImageSpec": Builder that returns a new image spec with force push enabled. """ new_image_spec = copy.deepcopy(self) + new_image_spec._is_force_push = True self._is_force_push = True return new_image_spec diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 0b2a286f94..b42be184f8 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="guqLm2fbPHgwk0EiXmXI9g.."), + default_image=Image(name="default", fqn="flytekit", tag="DgQMqIi61py4I4P5iOeS0Q.."), images=[ - Image(name="default", fqn="flytekit", tag="guqLm2fbPHgwk0EiXmXI9g.."), + 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 be27330ee3..961842838e 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -26,6 +26,7 @@ def test_image_spec(mock_image_spec_builder): requirements=REQUIREMENT_FILE, registry_config=REGISTRY_CONFIG_FILE, ) + assert image_spec._is_force_push is False image_spec = image_spec.with_commands("echo hello") image_spec = image_spec.with_packages("numpy") @@ -48,7 +49,7 @@ def test_image_spec(mock_image_spec_builder): assert image_spec.pip_index is None assert image_spec.is_container() is True assert image_spec.commands == ["echo hello"] - assert image_spec.is_force_push is True + assert image_spec._is_force_push is True tag = calculate_hash_from_image_spec(image_spec) assert image_spec.image_name() == f"flytekit:{tag}" From c7b4fcb6218f8cd331708bb671ba9c52cac54597 Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Mon, 4 Mar 2024 23:35:33 +0800 Subject: [PATCH 07/12] refactor: refactor variable names across files - Change the variable `is_force_push` to `_is_force_push` Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 28a3062222..d47eb9fe87 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -44,7 +44,7 @@ class ImageSpec: pip_index: Specify the custom pip index url registry_config: Specify the path to a JSON registry config file commands: Command to run during the building process - is_force_push: forcefully pushed to the registry + _is_force_push: forcefully pushed to the registry """ name: str = "flytekit" From d3a0fb5e20d1dad4d43f34939efafda920c598fb Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Wed, 6 Mar 2024 09:17:38 +0800 Subject: [PATCH 08/12] refactor: refactor class attributes in ImageSpec - Remove the `_is_force_push` attribute from `ImageSpec` class Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index d47eb9fe87..8f89af0191 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -44,7 +44,6 @@ class ImageSpec: pip_index: Specify the custom pip index url registry_config: Specify the path to a JSON registry config file commands: Command to run during the building process - _is_force_push: forcefully pushed to the registry """ name: str = "flytekit" @@ -189,7 +188,6 @@ def force_push(self) -> "ImageSpec": """ new_image_spec = copy.deepcopy(self) new_image_spec._is_force_push = True - self._is_force_push = True return new_image_spec From 747178516c67af65952beced49be1f811b751635 Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Wed, 6 Mar 2024 17:38:44 +0800 Subject: [PATCH 09/12] refactor: refactor image building and registration logic - Refactor the `ImageBuildEngine` class to improve readability and maintainability - Update the logic for handling image building and registration - Fix formatting and messaging inconsistencies in the code Signed-off-by: jason.lai --- flytekit/image_spec/image_spec.py | 44 +++++++++++++++++-------------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 8f89af0191..0c0d9ee01e 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -232,29 +232,33 @@ def build(cls, image_spec: ImageSpec) -> str: builder = image_spec.builder img_name = image_spec.image_name() - if (img_name in cls._BUILT_IMAGES or image_spec.exist()) and not image_spec._is_force_push: - click.secho(f"Image {img_name} found. Skip building.", fg="blue") - else: + if img_name in cls._BUILT_IMAGES or image_spec.exist(): if image_spec._is_force_push: click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") + cls.register_image(builder, image_spec, img_name) else: - click.secho(f"Image {img_name} not found. Building...", fg="blue") - if builder not in cls._REGISTRY: - raise Exception(f"Builder {builder} is not registered.") - if builder == "envd": - envd_version = metadata.version("envd") - # flytekit v1.10.2+ copies the workflow code to the WorkDir specified in the Dockerfile. However, envd<0.3.39 - # overwrites the WorkDir when building the image, resulting in a permission issue when flytekit downloads the file. - if Version(envd_version) < Version("0.3.39"): - raise FlyteAssertion( - f"envd version {envd_version} is not compatible with flytekit>v1.10.2." - f" Please upgrade envd to v0.3.39+." - ) - - fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) - if fully_qualified_image_name is not None: - cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name - cls._BUILT_IMAGES.add(img_name) + click.secho(f"Image {img_name} found. Skip building.", fg="blue") + else: + click.secho(f"Image {img_name} not found. building...", fg="blue") + cls.register_image(builder, image_spec, img_name) + + @classmethod + def register_image(cls, builder, image_spec, img_name): + if builder not in cls._REGISTRY: + raise Exception(f"Builder {builder} is not registered.") + if builder == "envd": + envd_version = metadata.version("envd") + # flytekit v1.10.2+ copies the workflow code to the WorkDir specified in the Dockerfile. However, envd<0.3.39 + # overwrites the WorkDir when building the image, resulting in a permission issue when flytekit downloads the file. + if Version(envd_version) < Version("0.3.39"): + raise FlyteAssertion( + f"envd version {envd_version} is not compatible with flytekit>v1.10.2." + f" Please upgrade envd to v0.3.39+." + ) + fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) + if fully_qualified_image_name is not None: + cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name + cls._BUILT_IMAGES.add(img_name) @lru_cache From 32d39c79c7dca19767d2dfd0332bff39a12e962f Mon Sep 17 00:00:00 2001 From: "jason.lai" Date: Thu, 7 Mar 2024 01:43:23 +0800 Subject: [PATCH 10/12] test: refactor image building and testing process - Add a new test for building an existing image with force push Signed-off-by: jason.lai --- .../flytekit/unit/core/image_spec/test_image_spec.py | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 961842838e..f7507e88e4 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -100,3 +100,14 @@ def test_image_spec_engine_priority(): assert image_spec.image_name() == new_image_name del ImageBuildEngine._REGISTRY["build_10"] del ImageBuildEngine._REGISTRY["build_default"] + + +def test_build_existing_image_with_force_push(): + image_spec = Mock() + image_spec.exist.return_value = True + image_spec._is_force_push = True + + ImageBuildEngine.register_image = Mock() + + ImageBuildEngine.build(image_spec) + ImageBuildEngine.register_image.assert_called_once() From 1cb8860220740cd86ad875833760a12a81626ae3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 9 Mar 2024 21:42:20 -0800 Subject: [PATCH 11/12] Add FLYTE_FORCE_PUSH_IMAGE_SPEC Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 13 ++++++------- .../unit/core/image_spec/test_image_spec.py | 4 ++-- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index c7372db713..326e8a0e80 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -18,6 +18,7 @@ DOCKER_HUB = "docker.io" _F_IMG_ID = "_F_IMG_ID" +FLYTE_FORCE_PUSH_IMAGE_SPEC = "FLYTE_FORCE_PUSH_IMAGE_SPEC" @dataclass @@ -67,7 +68,7 @@ class ImageSpec: def __post_init__(self): self.name = self.name.lower() - self._is_force_push = False # False by default + self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False) # False by default if self.registry: self.registry = self.registry.lower() @@ -213,7 +214,6 @@ class ImageBuildEngine: """ _REGISTRY: typing.Dict[str, Tuple[ImageSpecBuilder, int]] = {} - _BUILT_IMAGES: typing.Set[str] = set() # _IMAGE_NAME_TO_REAL_NAME is used to keep track of the fully qualified image name # returned by the image builder. This allows ImageSpec to map from `image_spc.image_name()` # to the real qualified name. @@ -232,18 +232,18 @@ def build(cls, image_spec: ImageSpec) -> str: builder = image_spec.builder img_name = image_spec.image_name() - if img_name in cls._BUILT_IMAGES or image_spec.exist(): + if image_spec.exist(): if image_spec._is_force_push: click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") - cls.register_image(builder, image_spec, img_name) + cls._build_image(builder, image_spec, img_name) else: click.secho(f"Image {img_name} found. Skip building.", fg="blue") else: click.secho(f"Image {img_name} not found. building...", fg="blue") - cls.register_image(builder, image_spec, img_name) + cls._build_image(builder, image_spec, img_name) @classmethod - def register_image(cls, builder, image_spec, img_name): + def _build_image(cls, builder, image_spec, img_name): if builder not in cls._REGISTRY: raise Exception(f"Builder {builder} is not registered.") if builder == "envd": @@ -258,7 +258,6 @@ def register_image(cls, builder, image_spec, img_name): fully_qualified_image_name = cls._REGISTRY[builder][0].build_image(image_spec) if fully_qualified_image_name is not None: cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name - cls._BUILT_IMAGES.add(img_name) @lru_cache 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 a68d3625e3..724bb83a61 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -108,7 +108,7 @@ def test_build_existing_image_with_force_push(): image_spec.exist.return_value = True image_spec._is_force_push = True - ImageBuildEngine.register_image = Mock() + ImageBuildEngine.build_image = Mock() ImageBuildEngine.build(image_spec) - ImageBuildEngine.register_image.assert_called_once() + ImageBuildEngine.build_image.assert_called_once() From 7a71b5754a6e7894c9d85d030d6cae686aa49247 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 9 Mar 2024 21:49:17 -0800 Subject: [PATCH 12/12] Fix tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/image_spec/test_image_spec.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 724bb83a61..f1b0bbe975 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -108,7 +108,7 @@ def test_build_existing_image_with_force_push(): image_spec.exist.return_value = True image_spec._is_force_push = True - ImageBuildEngine.build_image = Mock() + ImageBuildEngine._build_image = Mock() ImageBuildEngine.build(image_spec) - ImageBuildEngine.build_image.assert_called_once() + ImageBuildEngine._build_image.assert_called_once()