From 3966d1a0a1e33137a4bc41d9860d4ed5e264cbdf Mon Sep 17 00:00:00 2001 From: "Thomas J. Fan" Date: Wed, 24 Jan 2024 16:52:45 -0500 Subject: [PATCH] Extends ImageSpec to accept image names from plugin and have priority for plugins (#2119) * Extends image spec to have priority and accept image names from plugin Signed-off-by: Thomas J. Fan * TST Simplify implemenation of registry Signed-off-by: Thomas J. Fan * FIX Be more specific about max key Signed-off-by: Thomas J. Fan * FIX Fixes failing test Signed-off-by: Thomas J. Fan * DOC Adds comment about _IMAGE_NAME_TO_REAL_NAME Signed-off-by: Thomas J. Fan --------- Signed-off-by: Thomas J. Fan --- flytekit/image_spec/image_spec.py | 48 +++++++++++++------ .../flytekit-envd/tests/test_image_spec.py | 16 ++++++- .../unit/core/image_spec/test_image_spec.py | 22 +++++++++ 3 files changed, 71 insertions(+), 15 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 35001eb31c..7a8ef547da 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -8,7 +8,7 @@ from dataclasses import asdict, dataclass from functools import lru_cache from importlib import metadata -from typing import List, Optional, Union +from typing import Dict, List, Optional, Tuple, Union import click import requests @@ -48,7 +48,7 @@ class ImageSpec: name: str = "flytekit" python_version: str = None # Use default python in the base image if None. - builder: str = "envd" + builder: Optional[str] = None source_root: Optional[str] = None env: Optional[typing.Dict[str, str]] = None registry: Optional[str] = None @@ -71,9 +71,15 @@ def __post_init__(self): self.registry = self.registry.lower() def image_name(self) -> str: - """ - return full image name with tag. - """ + """Full image name with tag.""" + image_name = self._image_name() + try: + return ImageBuildEngine._IMAGE_NAME_TO_REAL_NAME[image_name] + except KeyError: + return image_name + + def _image_name(self) -> str: + """Construct full image name with tag.""" tag = calculate_hash_from_image_spec(self) container_image = f"{self.name}:{tag}" if self.registry: @@ -178,12 +184,15 @@ def with_apt_packages(self, apt_packages: Union[str, List[str]]) -> "ImageSpec": class ImageSpecBuilder: @abstractmethod - def build_image(self, image_spec: ImageSpec): + def build_image(self, image_spec: ImageSpec) -> Optional[str]: """ Build the docker image and push it to the registry. Args: image_spec: image spec of the task. + + Returns: + fully_qualified_image_name: Fully qualified image name. If None, then `image_spec.image_name()` is used. """ raise NotImplementedError("This method is not implemented in the base class.") @@ -193,24 +202,33 @@ class ImageBuildEngine: ImageBuildEngine contains a list of builders that can be used to build an ImageSpec. """ - _REGISTRY: typing.Dict[str, ImageSpecBuilder] = {} + _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. + _IMAGE_NAME_TO_REAL_NAME: Dict[str, str] = {} @classmethod - def register(cls, builder_type: str, image_spec_builder: ImageSpecBuilder): - cls._REGISTRY[builder_type] = image_spec_builder + def register(cls, builder_type: str, image_spec_builder: ImageSpecBuilder, priority: int = 5): + cls._REGISTRY[builder_type] = (image_spec_builder, priority) @classmethod @lru_cache - def build(cls, image_spec: ImageSpec): + def build(cls, image_spec: ImageSpec) -> str: + if image_spec.builder is None and cls._REGISTRY: + builder = max(cls._REGISTRY, key=lambda name: cls._REGISTRY[name][1]) + else: + builder = image_spec.builder + img_name = image_spec.image_name() if img_name in cls._BUILT_IMAGES or image_spec.exist(): 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.builder not in cls._REGISTRY: - raise Exception(f"Builder {image_spec.builder} is not registered.") - if image_spec.builder == "envd": + 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. @@ -220,7 +238,9 @@ def build(cls, image_spec: ImageSpec): f" Please upgrade envd to v0.3.39+." ) - cls._REGISTRY[image_spec.builder].build_image(image_spec) + 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) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 36adebd346..d77b2ca89b 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -1,9 +1,23 @@ from pathlib import Path from textwrap import dedent +import pytest from flytekitplugins.envd.image_builder import EnvdImageSpecBuilder, create_envd_config -from flytekit.image_spec.image_spec import ImageSpec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec + + +@pytest.fixture(scope="module", autouse=True) +def register_envd_higher_priority(): + # Register a new envd platform with the highest priority so the test in this file uses envd + highest_priority_builder = max(ImageBuildEngine._REGISTRY, key=ImageBuildEngine._REGISTRY.get) + highest_priority = ImageBuildEngine._REGISTRY[highest_priority_builder][1] + yield ImageBuildEngine.register( + "envd_high_priority", + EnvdImageSpecBuilder(), + priority=highest_priority + 1, + ) + del ImageBuildEngine._REGISTRY["envd_high_priority"] def test_image_spec(): 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 81151beddd..38727d02fd 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -1,4 +1,5 @@ import os +from unittest.mock import Mock import pytest @@ -75,3 +76,24 @@ def test_image_spec(mock_image_spec_builder): # ImageSpec should be immutable image_spec.with_commands("ls") assert image_spec.commands == ["echo hello"] + + +def test_image_spec_engine_priority(): + image_spec = ImageSpec(name="FLYTEKIT") + image_name = image_spec.image_name() + + new_image_name = f"fqn.xyz/{image_name}" + mock_image_builder_10 = Mock() + mock_image_builder_10.build_image.return_value = new_image_name + mock_image_builder_default = Mock() + mock_image_builder_default.build_image.side_effect = ValueError("should not be called") + + ImageBuildEngine.register("build_10", mock_image_builder_10, priority=10) + ImageBuildEngine.register("build_default", mock_image_builder_default) + + ImageBuildEngine.build(image_spec) + mock_image_builder_10.build_image.assert_called_once_with(image_spec) + + assert image_spec.image_name() == new_image_name + del ImageBuildEngine._REGISTRY["build_10"] + del ImageBuildEngine._REGISTRY["build_default"]