From f34ab17a7bf57e4eb9ab952b8b33393771fa7b25 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 15:36:24 -0700 Subject: [PATCH 01/23] wip Signed-off-by: Kevin Su --- flytekit/tools/fast_registration.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flytekit/tools/fast_registration.py b/flytekit/tools/fast_registration.py index e596e62f38..005d266051 100644 --- a/flytekit/tools/fast_registration.py +++ b/flytekit/tools/fast_registration.py @@ -82,6 +82,8 @@ def compute_digest(source: os.PathLike, filter: Optional[callable] = None) -> st def _filehash_update(path: os.PathLike, hasher: hashlib._Hash) -> None: + if not os.path.exists(path): + return None blocksize = 65536 with open(path, "rb") as f: bytes = f.read(blocksize) From 162f475769a04e713681faf9555e9c78e6feac3b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 18:33:42 -0700 Subject: [PATCH 02/23] test Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 5 +++-- flytekit/core/python_auto_container.py | 12 ++++++++---- flytekit/remote/remote.py | 2 +- flytekit/tools/fast_registration.py | 2 -- flytekit/tools/serialize_helpers.py | 1 + flytekit/tools/translator.py | 9 +++++++-- 6 files changed, 20 insertions(+), 11 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index b5278be053..48361902c0 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -553,8 +553,9 @@ def _run(*args, **kwargs): image_config = run_level_params.image_config image_config = patch_image_config(config_file, image_config) - - with context_manager.FlyteContextManager.with_context(remote.context.new_builder()): + print(context_manager.FlyteEntities.entities) + with context_manager.FlyteContextManager.with_context(remote.context.new_builder()) as nctx: + # breakpoint() remote_entity = remote.register_script( entity, project=run_level_params.project, diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 29167ac031..bcbc1bca82 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -180,7 +180,7 @@ def get_image(self, settings: SerializationSettings) -> str: if isinstance(self.container_image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config) + return get_registerable_container_image(self.container_image, settings.image_config, self.name) def get_container(self, settings: SerializationSettings) -> _task_model.Container: # if pod_template is not None, return None here but in get_k8s_pod, return pod_template merged with container @@ -264,7 +264,7 @@ def get_all_tasks(self) -> List[PythonAutoContainerTask]: # type: ignore default_task_resolver = DefaultTaskResolver() -def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: ImageConfig) -> str: +def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: str = Optional[None]) -> str: """ Resolve the image to the real image name that should be used for registration. 1. If img is a ImageSpec, it will be built and the image name will be returned @@ -273,11 +273,15 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :param img: Configured image or image spec :param cfg: Registration configuration + :param task_name: The name of the container task. :return: """ if isinstance(img, ImageSpec): - ImageBuildEngine.build(img) - return img.image_name() + image_name = cfg.find_image(task_name) if task_name else None + if not image_name: + ImageBuildEngine.build(img) + image_name = img.image_name() + return image_name if img is not None and img != "": matches = _IMAGE_REPLACE_REGEX.findall(img) diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index 584501f137..fcbfe9b251 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -33,7 +33,7 @@ from flytekit.clients.friendly import SynchronousFlyteClient from flytekit.clients.helpers import iterate_node_executions, iterate_task_executions from flytekit.configuration import Config, FastSerializationSettings, ImageConfig, SerializationSettings -from flytekit.core import constants, utils +from flytekit.core import constants, utils, context_manager from flytekit.core.artifact import Artifact from flytekit.core.base_task import PythonTask from flytekit.core.context_manager import FlyteContext, FlyteContextManager diff --git a/flytekit/tools/fast_registration.py b/flytekit/tools/fast_registration.py index 005d266051..e596e62f38 100644 --- a/flytekit/tools/fast_registration.py +++ b/flytekit/tools/fast_registration.py @@ -82,8 +82,6 @@ def compute_digest(source: os.PathLike, filter: Optional[callable] = None) -> st def _filehash_update(path: os.PathLike, hasher: hashlib._Hash) -> None: - if not os.path.exists(path): - return None blocksize = 65536 with open(path, "rb") as f: bytes = f.read(blocksize) diff --git a/flytekit/tools/serialize_helpers.py b/flytekit/tools/serialize_helpers.py index 8d4cfcb99c..a09afedd69 100644 --- a/flytekit/tools/serialize_helpers.py +++ b/flytekit/tools/serialize_helpers.py @@ -53,6 +53,7 @@ def get_registrable_entities( # TODO: Clean up the copy() - it's here because we call get_default_launch_plan, which may create a LaunchPlan # object, which gets added to the FlyteEntities.entities list, which we're iterating over. for entity in flyte_context.FlyteEntities.entities.copy(): + # breakpoint() if isinstance(entity, PythonTask) or isinstance(entity, WorkflowBase) or isinstance(entity, LaunchPlan): get_serializable(new_api_serializable_entities, ctx.serialization_settings, entity, options=options) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index b49639d23a..e033623108 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -7,8 +7,8 @@ from flyteidl.admin import schedule_pb2 from flytekit import PythonFunctionTask, SourceCode -from flytekit.configuration import SerializationSettings -from flytekit.core import constants as _common_constants +from flytekit.configuration import SerializationSettings, Image +from flytekit.core import constants as _common_constants, context_manager from flytekit.core.array_node_map_task import ArrayNodeMapTask from flytekit.core.base_task import PythonTask from flytekit.core.condition import BranchNode @@ -176,6 +176,11 @@ def get_serializable_task( ) if isinstance(entity, PythonFunctionTask) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC: + print(len(context_manager.FlyteEntities.entities)) + for e in context_manager.FlyteEntities.entities: + if isinstance(e, PythonAutoContainerTask): + settings.image_config.images.append(Image.look_up_image_info(e.name, e.get_image(settings))) + # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state # from the serialization context. This is passed through an environment variable, that is read from # during dynamic serialization From b4bef35555ed392cb3a0ea61a83a21f8a7939bbf Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 18:40:40 -0700 Subject: [PATCH 03/23] test Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index bcbc1bca82..00ed2e8c57 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -277,7 +277,8 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :return: """ if isinstance(img, ImageSpec): - image_name = cfg.find_image(task_name) if task_name else None + image = cfg.find_image(task_name) if task_name else None + image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) image_name = img.image_name() From f38163b7a39f1e16351d04817bb3f8a9635ef72c Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 18:44:36 -0700 Subject: [PATCH 04/23] test Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 2 +- flytekit/tools/translator.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 00ed2e8c57..973d512f96 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -277,7 +277,7 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :return: """ if isinstance(img, ImageSpec): - image = cfg.find_image(task_name) if task_name else None + image = cfg.find_image(f"ft_{task_name}") if task_name else None image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index e033623108..f29f4f4fe0 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -179,7 +179,7 @@ def get_serializable_task( print(len(context_manager.FlyteEntities.entities)) for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): - settings.image_config.images.append(Image.look_up_image_info(e.name, e.get_image(settings))) + settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state # from the serialization context. This is passed through an environment variable, that is read from From f44a2550f7315f52810b77bb4df28e6d5ee5fd62 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 18:53:16 -0700 Subject: [PATCH 05/23] test Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 2 ++ flytekit/tools/translator.py | 1 + 2 files changed, 3 insertions(+) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 973d512f96..696f1dbc95 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -277,6 +277,8 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :return: """ if isinstance(img, ImageSpec): + print(cfg.images) + print(f"ft_{task_name}") image = cfg.find_image(f"ft_{task_name}") if task_name else None image_name = image.full if image else None if not image_name: diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index f29f4f4fe0..b56e1abbf9 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -179,6 +179,7 @@ def get_serializable_task( print(len(context_manager.FlyteEntities.entities)) for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): + print(f"ft_{e.name}") settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state From 7bdd483056b86b1c7f7a1452e2d0a6d061f2d1a3 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 18:59:47 -0700 Subject: [PATCH 06/23] test Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 1 + flytekit/tools/translator.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 696f1dbc95..371d26966b 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -276,6 +276,7 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :param task_name: The name of the container task. :return: """ + print("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") if isinstance(img, ImageSpec): print(cfg.images) print(f"ft_{task_name}") diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index b56e1abbf9..6d271a13b0 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -179,7 +179,7 @@ def get_serializable_task( print(len(context_manager.FlyteEntities.entities)) for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): - print(f"ft_{e.name}") + print("aaaaaa", f"ft_{e.name}") settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state From 2918cb930286edc88f2f3cedfc0bf63ca3bc1310 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 19:18:06 -0700 Subject: [PATCH 07/23] test Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 4 +-- flytekit/core/python_auto_container.py | 3 --- flytekit/tools/translator.py | 1 - .../flytekit/unit/core/test_serialization.py | 26 ++++++++++++++++--- 4 files changed, 23 insertions(+), 11 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index 48361902c0..e843b33a6a 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -553,9 +553,7 @@ def _run(*args, **kwargs): image_config = run_level_params.image_config image_config = patch_image_config(config_file, image_config) - print(context_manager.FlyteEntities.entities) - with context_manager.FlyteContextManager.with_context(remote.context.new_builder()) as nctx: - # breakpoint() + with context_manager.FlyteContextManager.with_context(remote.context.new_builder()): remote_entity = remote.register_script( entity, project=run_level_params.project, diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 371d26966b..973d512f96 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -276,10 +276,7 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :param task_name: The name of the container task. :return: """ - print("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa") if isinstance(img, ImageSpec): - print(cfg.images) - print(f"ft_{task_name}") image = cfg.find_image(f"ft_{task_name}") if task_name else None image_name = image.full if image else None if not image_name: diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 6d271a13b0..f29f4f4fe0 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -179,7 +179,6 @@ def get_serializable_task( print(len(context_manager.FlyteEntities.entities)) for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): - print("aaaaaa", f"ft_{e.name}") settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 9b11a2a16a..8c7f000e4a 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -6,12 +6,13 @@ import pytest import flytekit.configuration -from flytekit import ContainerTask, kwtypes +from flytekit import ContainerTask, kwtypes, ImageSpec from flytekit.configuration import Image, ImageConfig, SerializationSettings from flytekit.core.condition import conditional from flytekit.core.python_auto_container import get_registerable_container_image from flytekit.core.task import task from flytekit.core.workflow import workflow +from flytekit.image_spec.image_spec import ImageBuildEngine from flytekit.models.admin.workflow import WorkflowSpec from flytekit.models.types import SimpleType from flytekit.tools.translator import get_serializable @@ -250,7 +251,9 @@ def test_bad_configuration(): get_registerable_container_image(container_image, image_config) -def test_serialization_images(): +def test_serialization_images(mock_image_spec_builder): + ImageBuildEngine.register("test", mock_image_spec_builder) + @task(container_image="{{.image.xyz.fqn}}:{{.image.xyz.version}}") def t1(a: int) -> int: return a @@ -271,10 +274,22 @@ def t5(a: int) -> int: def t6(a: int) -> int: return a + image_spec = ImageSpec( + packages=["mypy"], + apt_packages=["git"], + registry="ghcr.io/flyteorg", + builder="test", + ) + + @task(container_image=image_spec) + def t7(a: int) -> int: + return a + with mock.patch.dict(os.environ, {"FLYTE_INTERNAL_IMAGE": "docker.io/default:version"}): imgs = ImageConfig.auto( config_file=os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/images.config") ) + imgs.images.append(Image(name=f"ft_{t7.name}", fqn="docker.io/t7", tag="latest")) rs = flytekit.configuration.SerializationSettings( project="project", domain="domain", @@ -295,8 +310,11 @@ def t6(a: int) -> int: t5_spec = get_serializable(OrderedDict(), rs, t5) assert t5_spec.template.container.image == "docker.io/org/myimage:latest" - t5_spec = get_serializable(OrderedDict(), rs, t6) - assert t5_spec.template.container.image == "docker.io/xyz_123:v1" + t6_spec = get_serializable(OrderedDict(), rs, t6) + assert t6_spec.template.container.image == "docker.io/xyz_123:v1" + + t7_spec = get_serializable(OrderedDict(), rs, t7) + assert t7_spec.template.container.image == "docker.io/t7:latest" def test_serialization_command1(): From 46671be977e3d9ea764fe9136ce00880cc321ec1 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 19:27:24 -0700 Subject: [PATCH 08/23] test Signed-off-by: Kevin Su --- flytekit/clis/sdk_in_container/run.py | 1 + flytekit/core/container_task.py | 2 +- flytekit/core/python_customized_container_task.py | 2 +- flytekit/tools/serialize_helpers.py | 1 - flytekit/tools/translator.py | 1 - plugins/flytekit-spark/flytekitplugins/spark/task.py | 2 +- tests/flytekit/unit/core/test_python_auto_container.py | 6 +++++- 7 files changed, 9 insertions(+), 6 deletions(-) diff --git a/flytekit/clis/sdk_in_container/run.py b/flytekit/clis/sdk_in_container/run.py index e843b33a6a..b5278be053 100644 --- a/flytekit/clis/sdk_in_container/run.py +++ b/flytekit/clis/sdk_in_container/run.py @@ -553,6 +553,7 @@ def _run(*args, **kwargs): image_config = run_level_params.image_config image_config = patch_image_config(config_file, image_config) + with context_manager.FlyteContextManager.with_context(remote.context.new_builder()): remote_entity = remote.register_script( entity, diff --git a/flytekit/core/container_task.py b/flytekit/core/container_task.py index 66fe522c07..5d2ab09c67 100644 --- a/flytekit/core/container_task.py +++ b/flytekit/core/container_task.py @@ -276,7 +276,7 @@ def _get_image(self, settings: SerializationSettings) -> str: if isinstance(self._image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self._image.source_root = settings.source_root - return get_registerable_container_image(self._image, settings.image_config) + return get_registerable_container_image(self._image, settings.image_config, self.name) def _get_container(self, settings: SerializationSettings) -> _task_model.Container: env = settings.env or {} diff --git a/flytekit/core/python_customized_container_task.py b/flytekit/core/python_customized_container_task.py index fd3ab4a8f4..29a8505139 100644 --- a/flytekit/core/python_customized_container_task.py +++ b/flytekit/core/python_customized_container_task.py @@ -164,7 +164,7 @@ def get_image(self, settings: SerializationSettings) -> str: if isinstance(self.container_image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config) + return get_registerable_container_image(self.container_image, settings.image_config, self.name) def get_container(self, settings: SerializationSettings) -> _task_model.Container: env = {**settings.env, **self.environment} if self.environment else settings.env diff --git a/flytekit/tools/serialize_helpers.py b/flytekit/tools/serialize_helpers.py index a09afedd69..8d4cfcb99c 100644 --- a/flytekit/tools/serialize_helpers.py +++ b/flytekit/tools/serialize_helpers.py @@ -53,7 +53,6 @@ def get_registrable_entities( # TODO: Clean up the copy() - it's here because we call get_default_launch_plan, which may create a LaunchPlan # object, which gets added to the FlyteEntities.entities list, which we're iterating over. for entity in flyte_context.FlyteEntities.entities.copy(): - # breakpoint() if isinstance(entity, PythonTask) or isinstance(entity, WorkflowBase) or isinstance(entity, LaunchPlan): get_serializable(new_api_serializable_entities, ctx.serialization_settings, entity, options=options) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index f29f4f4fe0..99c69e3ffe 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -176,7 +176,6 @@ def get_serializable_task( ) if isinstance(entity, PythonFunctionTask) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC: - print(len(context_manager.FlyteEntities.entities)) for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index 8a8c3b2b5b..82f655fc78 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -140,7 +140,7 @@ def get_image(self, settings: SerializationSettings) -> str: # Ensure that the code is always copied into the image, even during fast-registration. self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config) + return get_registerable_container_image(self.container_image, settings.image_config, self.name) def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: job = SparkJob( diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index 58492fca06..ab3c19eeb3 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -55,9 +55,13 @@ def serialization_settings(request): def test_image_name_interpolation(default_image_config): + new_img_cfg = ImageConfig.create_from(default_image_config.default_image, other_images=[Image.look_up_image_info("ft_d1", "flyte/test:d1")]) img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" - img = get_registerable_container_image(img=img_to_interpolate, cfg=default_image_config) + img = get_registerable_container_image(img=img_to_interpolate, cfg=new_img_cfg) assert img == "docker.io/xyz:some-git-hash-special" + image = ImageSpec(name="image-1", registry="localhost:30000", builder="test") + img = get_registerable_container_image(img=image, cfg=new_img_cfg, task_name="d1") + assert img == "flyte/test:d1" class DummyAutoContainerTask(PythonAutoContainerTask): From 70a646c363da82faf88dd399fccbeb1a9c6b86a9 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 19:27:51 -0700 Subject: [PATCH 09/23] make fmt Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 4 +++- flytekit/remote/remote.py | 2 +- flytekit/tools/translator.py | 5 +++-- tests/flytekit/unit/core/test_python_auto_container.py | 4 +++- tests/flytekit/unit/core/test_serialization.py | 2 +- 5 files changed, 11 insertions(+), 6 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 973d512f96..f5409346f1 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -264,7 +264,9 @@ def get_all_tasks(self) -> List[PythonAutoContainerTask]: # type: ignore default_task_resolver = DefaultTaskResolver() -def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: str = Optional[None]) -> str: +def get_registerable_container_image( + img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: str = Optional[None] +) -> str: """ Resolve the image to the real image name that should be used for registration. 1. If img is a ImageSpec, it will be built and the image name will be returned diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py index fcbfe9b251..584501f137 100644 --- a/flytekit/remote/remote.py +++ b/flytekit/remote/remote.py @@ -33,7 +33,7 @@ from flytekit.clients.friendly import SynchronousFlyteClient from flytekit.clients.helpers import iterate_node_executions, iterate_task_executions from flytekit.configuration import Config, FastSerializationSettings, ImageConfig, SerializationSettings -from flytekit.core import constants, utils, context_manager +from flytekit.core import constants, utils from flytekit.core.artifact import Artifact from flytekit.core.base_task import PythonTask from flytekit.core.context_manager import FlyteContext, FlyteContextManager diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 99c69e3ffe..54dab6c1ed 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -7,8 +7,9 @@ from flyteidl.admin import schedule_pb2 from flytekit import PythonFunctionTask, SourceCode -from flytekit.configuration import SerializationSettings, Image -from flytekit.core import constants as _common_constants, context_manager +from flytekit.configuration import Image, SerializationSettings +from flytekit.core import constants as _common_constants +from flytekit.core import context_manager from flytekit.core.array_node_map_task import ArrayNodeMapTask from flytekit.core.base_task import PythonTask from flytekit.core.condition import BranchNode diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index ab3c19eeb3..4db8906fae 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -55,7 +55,9 @@ def serialization_settings(request): def test_image_name_interpolation(default_image_config): - new_img_cfg = ImageConfig.create_from(default_image_config.default_image, other_images=[Image.look_up_image_info("ft_d1", "flyte/test:d1")]) + new_img_cfg = ImageConfig.create_from( + default_image_config.default_image, other_images=[Image.look_up_image_info("ft_d1", "flyte/test:d1")] + ) img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" img = get_registerable_container_image(img=img_to_interpolate, cfg=new_img_cfg) assert img == "docker.io/xyz:some-git-hash-special" diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 8c7f000e4a..9bfbe8435d 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -6,7 +6,7 @@ import pytest import flytekit.configuration -from flytekit import ContainerTask, kwtypes, ImageSpec +from flytekit import ContainerTask, ImageSpec, kwtypes from flytekit.configuration import Image, ImageConfig, SerializationSettings from flytekit.core.condition import conditional from flytekit.core.python_auto_container import get_registerable_container_image From 28d4da3a49c98a9f8dc93fef0945e22ca3cd208a Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 19:35:30 -0700 Subject: [PATCH 10/23] test Signed-off-by: Kevin Su --- flytekit/tools/translator.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 54dab6c1ed..dc9c292ad2 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -6,7 +6,7 @@ from flyteidl.admin import schedule_pb2 -from flytekit import PythonFunctionTask, SourceCode +from flytekit import PythonFunctionTask, SourceCode, ImageSpec from flytekit.configuration import Image, SerializationSettings from flytekit.core import constants as _common_constants from flytekit.core import context_manager @@ -179,7 +179,10 @@ def get_serializable_task( if isinstance(entity, PythonFunctionTask) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC: for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): - settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) + # 1. Build the ImageSpec for the entities that are inside the dynamic task, + # 2. Add images to the serialization context, so the dynamic task can look it up at runtime. + if isinstance(e.container_image, ImageSpec): + settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state # from the serialization context. This is passed through an environment variable, that is read from From 11595a6c6bf814a5057424ad733763d8e1437ad0 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 19:58:46 -0700 Subject: [PATCH 11/23] None check Signed-off-by: Kevin Su --- flytekit/tools/translator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index dc9c292ad2..c2e74ec60d 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -6,8 +6,8 @@ from flyteidl.admin import schedule_pb2 -from flytekit import PythonFunctionTask, SourceCode, ImageSpec -from flytekit.configuration import Image, SerializationSettings +from flytekit import ImageSpec, PythonFunctionTask, SourceCode +from flytekit.configuration import Image, ImageConfig, SerializationSettings from flytekit.core import constants as _common_constants from flytekit.core import context_manager from flytekit.core.array_node_map_task import ArrayNodeMapTask @@ -182,6 +182,8 @@ def get_serializable_task( # 1. Build the ImageSpec for the entities that are inside the dynamic task, # 2. Add images to the serialization context, so the dynamic task can look it up at runtime. if isinstance(e.container_image, ImageSpec): + if settings.image_config.images is None: + settings.image_config = ImageConfig.create_from(settings.image_config.default_image) settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state From 375db5cebde2c03b4cd350ff994e6e60b1d599eb Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 20:00:46 -0700 Subject: [PATCH 12/23] lint Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index f5409346f1..d783106e87 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -265,7 +265,7 @@ def get_all_tasks(self) -> List[PythonAutoContainerTask]: # type: ignore def get_registerable_container_image( - img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: str = Optional[None] + img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: Optional[str] = None ) -> str: """ Resolve the image to the real image name that should be used for registration. From 4aa5d2e98d7c76e91eea3981d7248085bace98b7 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 2 Jun 2024 20:16:10 -0700 Subject: [PATCH 13/23] test Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_node_creation.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py index fc3284ca10..684f49031b 100644 --- a/tests/flytekit/unit/core/test_node_creation.py +++ b/tests/flytekit/unit/core/test_node_creation.py @@ -14,12 +14,15 @@ from flytekit.core.workflow import workflow from flytekit.exceptions.user import FlyteAssertion from flytekit.extras.accelerators import A100, T4 +from flytekit.image_spec.image_spec import ImageBuildEngine from flytekit.models import literals as _literal_models from flytekit.models.task import Resources as _resources_models from flytekit.tools.translator import get_serializable -def test_normal_task(): +def test_normal_task(mock_image_spec_builder): + ImageBuildEngine.register("test", mock_image_spec_builder) + @task def t1(a: str) -> str: return a + " world" From f00968b9d37c24a25de9ec104b68d8bd38b8909d Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 16:40:10 -0700 Subject: [PATCH 14/23] Use imageSpec name as key Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 2 +- flytekit/image_spec/image_spec.py | 7 ++++++- flytekit/tools/translator.py | 6 ++++-- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index d783106e87..aa6b110c05 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -279,7 +279,7 @@ def get_registerable_container_image( :return: """ if isinstance(img, ImageSpec): - image = cfg.find_image(f"ft_{task_name}") if task_name else None + image = cfg.find_image(f"ft_{img.lhs}") if task_name else None image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 98f6c05cdc..594d506a03 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -14,6 +14,7 @@ import requests from packaging.version import Version +from flytekit.core.tracker import TrackedInstance from flytekit.exceptions.user import FlyteAssertion DOCKER_HUB = "docker.io" @@ -22,7 +23,7 @@ @dataclass -class ImageSpec: +class ImageSpec(TrackedInstance): """ This class is used to specify the docker image that will be used to run the task. @@ -72,6 +73,9 @@ class ImageSpec: tag_format: Optional[str] = None def __post_init__(self): + self._instantiated_in = None + self._module_file = None + self._lhs = None self.name = self.name.lower() self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False) # False by default if self.registry: @@ -252,6 +256,7 @@ def build(cls, image_spec: ImageSpec): builder = image_spec.builder img_name = image_spec.image_name() + # print(img_name) if image_spec.exist(): if image_spec._is_force_push: click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index c2e74ec60d..32dee5c5d9 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -179,12 +179,14 @@ def get_serializable_task( if isinstance(entity, PythonFunctionTask) and entity.execution_mode == PythonFunctionTask.ExecutionBehavior.DYNAMIC: for e in context_manager.FlyteEntities.entities: if isinstance(e, PythonAutoContainerTask): - # 1. Build the ImageSpec for the entities that are inside the dynamic task, + # 1. Build the ImageSpec for all the entities that are inside the current context, # 2. Add images to the serialization context, so the dynamic task can look it up at runtime. if isinstance(e.container_image, ImageSpec): if settings.image_config.images is None: settings.image_config = ImageConfig.create_from(settings.image_config.default_image) - settings.image_config.images.append(Image.look_up_image_info(f"ft_{e.name}", e.get_image(settings))) + settings.image_config.images.append( + Image.look_up_image_info(f"ft_{e.container_image.lhs}", e.get_image(settings)) + ) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state # from the serialization context. This is passed through an environment variable, that is read from From 53fd3a490d08cea7a8df0a7b8032ca00c2221aae Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 16:47:07 -0700 Subject: [PATCH 15/23] fix tests Signed-off-by: Kevin Su --- flytekit/core/container_task.py | 2 +- flytekit/core/python_auto_container.py | 8 +++----- flytekit/core/python_customized_container_task.py | 2 +- flytekit/image_spec/image_spec.py | 1 - .../flytekit-spark/flytekitplugins/spark/task.py | 2 +- .../unit/core/test_python_auto_container.py | 9 ++++++--- tests/flytekit/unit/core/test_serialization.py | 15 +++++++-------- 7 files changed, 19 insertions(+), 20 deletions(-) diff --git a/flytekit/core/container_task.py b/flytekit/core/container_task.py index 5d2ab09c67..66fe522c07 100644 --- a/flytekit/core/container_task.py +++ b/flytekit/core/container_task.py @@ -276,7 +276,7 @@ def _get_image(self, settings: SerializationSettings) -> str: if isinstance(self._image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self._image.source_root = settings.source_root - return get_registerable_container_image(self._image, settings.image_config, self.name) + return get_registerable_container_image(self._image, settings.image_config) def _get_container(self, settings: SerializationSettings) -> _task_model.Container: env = settings.env or {} diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index aa6b110c05..caf994af1a 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -180,7 +180,7 @@ def get_image(self, settings: SerializationSettings) -> str: if isinstance(self.container_image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config, self.name) + return get_registerable_container_image(self.container_image, settings.image_config) def get_container(self, settings: SerializationSettings) -> _task_model.Container: # if pod_template is not None, return None here but in get_k8s_pod, return pod_template merged with container @@ -264,9 +264,7 @@ def get_all_tasks(self) -> List[PythonAutoContainerTask]: # type: ignore default_task_resolver = DefaultTaskResolver() -def get_registerable_container_image( - img: Optional[Union[str, ImageSpec]], cfg: ImageConfig, task_name: Optional[str] = None -) -> str: +def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: ImageConfig) -> str: """ Resolve the image to the real image name that should be used for registration. 1. If img is a ImageSpec, it will be built and the image name will be returned @@ -279,7 +277,7 @@ def get_registerable_container_image( :return: """ if isinstance(img, ImageSpec): - image = cfg.find_image(f"ft_{img.lhs}") if task_name else None + image = cfg.find_image(f"ft_{img.lhs}") image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) diff --git a/flytekit/core/python_customized_container_task.py b/flytekit/core/python_customized_container_task.py index 29a8505139..fd3ab4a8f4 100644 --- a/flytekit/core/python_customized_container_task.py +++ b/flytekit/core/python_customized_container_task.py @@ -164,7 +164,7 @@ def get_image(self, settings: SerializationSettings) -> str: if isinstance(self.container_image, ImageSpec): # Set the source root for the image spec if it's non-fast registration self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config, self.name) + return get_registerable_container_image(self.container_image, settings.image_config) def get_container(self, settings: SerializationSettings) -> _task_model.Container: env = {**settings.env, **self.environment} if self.environment else settings.env diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 594d506a03..83d67ae857 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -256,7 +256,6 @@ def build(cls, image_spec: ImageSpec): builder = image_spec.builder img_name = image_spec.image_name() - # print(img_name) if image_spec.exist(): if image_spec._is_force_push: click.secho(f"Image {img_name} found. but overwriting existing image.", fg="blue") diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py index 82f655fc78..8a8c3b2b5b 100644 --- a/plugins/flytekit-spark/flytekitplugins/spark/task.py +++ b/plugins/flytekit-spark/flytekitplugins/spark/task.py @@ -140,7 +140,7 @@ def get_image(self, settings: SerializationSettings) -> str: # Ensure that the code is always copied into the image, even during fast-registration. self.container_image.source_root = settings.source_root - return get_registerable_container_image(self.container_image, settings.image_config, self.name) + return get_registerable_container_image(self.container_image, settings.image_config) def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]: job = SparkJob( diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index 4db8906fae..142e4f86f4 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -54,15 +54,18 @@ def serialization_settings(request): return request.getfixturevalue(request.param) +image_spec = ImageSpec(name="image-1", registry="localhost:30000", builder="test") + + def test_image_name_interpolation(default_image_config): new_img_cfg = ImageConfig.create_from( - default_image_config.default_image, other_images=[Image.look_up_image_info("ft_d1", "flyte/test:d1")] + default_image_config.default_image, + other_images=[Image.look_up_image_info(f"ft_{image_spec.lhs}", "flyte/test:d1")], ) img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" img = get_registerable_container_image(img=img_to_interpolate, cfg=new_img_cfg) assert img == "docker.io/xyz:some-git-hash-special" - image = ImageSpec(name="image-1", registry="localhost:30000", builder="test") - img = get_registerable_container_image(img=image, cfg=new_img_cfg, task_name="d1") + img = get_registerable_container_image(img=image_spec, cfg=new_img_cfg) assert img == "flyte/test:d1" diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 9bfbe8435d..9805aea56c 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -26,6 +26,12 @@ env=None, image_config=ImageConfig(default_image=default_img, images=[default_img]), ) +image_spec = ImageSpec( + packages=["mypy"], + apt_packages=["git"], + registry="ghcr.io/flyteorg", + builder="test", +) def test_serialization(): @@ -274,13 +280,6 @@ def t5(a: int) -> int: def t6(a: int) -> int: return a - image_spec = ImageSpec( - packages=["mypy"], - apt_packages=["git"], - registry="ghcr.io/flyteorg", - builder="test", - ) - @task(container_image=image_spec) def t7(a: int) -> int: return a @@ -289,7 +288,7 @@ def t7(a: int) -> int: imgs = ImageConfig.auto( config_file=os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/images.config") ) - imgs.images.append(Image(name=f"ft_{t7.name}", fqn="docker.io/t7", tag="latest")) + imgs.images.append(Image(name=f"ft_{image_spec.lhs}", fqn="docker.io/t7", tag="latest")) rs = flytekit.configuration.SerializationSettings( project="project", domain="domain", From f8bc9d6c29d6931d9b5000e9e8ecf07d4df12382 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 17:07:01 -0700 Subject: [PATCH 16/23] fix tests Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 1 - .../flytekit/unit/core/test_container_task.py | 33 ++++++++++--------- .../unit/core/test_python_function_task.py | 4 ++- 3 files changed, 21 insertions(+), 17 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index caf994af1a..60a6522ce8 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -273,7 +273,6 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :param img: Configured image or image spec :param cfg: Registration configuration - :param task_name: The name of the container task. :return: """ if isinstance(img, ImageSpec): diff --git a/tests/flytekit/unit/core/test_container_task.py b/tests/flytekit/unit/core/test_container_task.py index 1281a9ec14..66d0ef1999 100644 --- a/tests/flytekit/unit/core/test_container_task.py +++ b/tests/flytekit/unit/core/test_container_task.py @@ -158,9 +158,11 @@ def test_pod_template(): assert serialized_pod_spec["runtimeClassName"] == "nvidia" +image_spec = ImageSpec(registry="flyte", base_image="r-base", builder="test-raw-container") + + def test_raw_container_with_image_spec(mock_image_spec_builder): ImageBuildEngine.register("test-raw-container", mock_image_spec_builder) - image_spec = ImageSpec(registry="flyte", base_image="r-base", builder="test-raw-container") calculate_ellipse_area_r = ContainerTask( name="ellipse-area-metadata-r", @@ -186,6 +188,21 @@ def test_raw_container_with_image_spec(mock_image_spec_builder): assert container.image == image_spec.image_name() +image_spec_1 = ImageSpec( + name="image-1", + packages=["numpy"], + registry="localhost:30000", + builder="test", +) + +image_spec_2 = ImageSpec( + name="image-2", + packages=["pandas"], + registry="localhost:30000", + builder="test", +) + + def test_container_task_image_spec(mock_image_spec_builder): default_image = Image(name="default", fqn="docker.io/xyz", tag="some-git-hash") default_image_config = ImageConfig(default_image=default_image) @@ -194,20 +211,6 @@ def test_container_task_image_spec(mock_image_spec_builder): project="p", domain="d", version="v", image_config=default_image_config, env={"FOO": "bar"} ) - image_spec_1 = ImageSpec( - name="image-1", - packages=["numpy"], - registry="localhost:30000", - builder="test", - ) - - image_spec_2 = ImageSpec( - name="image-2", - packages=["pandas"], - registry="localhost:30000", - builder="test", - ) - ps = V1PodSpec( containers=[ V1Container( diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 8ba875b664..a268ecb51c 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -36,6 +36,9 @@ def test_istestfunction(): assert istestfunction(tasks.tasks) is False +image_spec = ImageSpec(builder="test", python_version="3.7", registry="") + + def test_container_image_conversion(mock_image_spec_builder): default_img = Image(name="default", fqn="xyz.com/abc", tag="tag1") other_img = Image(name="other", fqn="xyz.com/other", tag="tag-other") @@ -111,7 +114,6 @@ def test_container_image_conversion(mock_image_spec_builder): ) ImageBuildEngine.register("test", mock_image_spec_builder) - image_spec = ImageSpec(builder="test", python_version="3.7", registry="") assert get_registerable_container_image(image_spec, cfg) == image_spec.image_name() From 77b760dd5b84e1693c7216cecd09a1e63a865149 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 17:16:22 -0700 Subject: [PATCH 17/23] fix tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_dynamic.py | 2 +- tests/flytekit/unit/core/test_node_creation.py | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/flytekit/unit/core/test_dynamic.py b/tests/flytekit/unit/core/test_dynamic.py index 7964548674..e49e09fad6 100644 --- a/tests/flytekit/unit/core/test_dynamic.py +++ b/tests/flytekit/unit/core/test_dynamic.py @@ -287,7 +287,7 @@ def dt(mode: int) -> int: raise ValueError("Invalid mode") entity_mapping = OrderedDict() - get_serializable_task(entity_mapping, settings, dt) + get_serializable_task(entity_mapping, settings.with_serialized_context(), dt) serialised_entities_iterator = iter(entity_mapping.values()) assert "t1" in next(serialised_entities_iterator).template.id.name diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py index 684f49031b..d5fe013e0c 100644 --- a/tests/flytekit/unit/core/test_node_creation.py +++ b/tests/flytekit/unit/core/test_node_creation.py @@ -20,9 +20,7 @@ from flytekit.tools.translator import get_serializable -def test_normal_task(mock_image_spec_builder): - ImageBuildEngine.register("test", mock_image_spec_builder) - +def test_normal_task(): @task def t1(a: str) -> str: return a + " world" From dae267d1458218fd5d419cd854ade4287a28ec7a Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 17:32:58 -0700 Subject: [PATCH 18/23] fix tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_node_creation.py | 1 - tests/flytekit/unit/remote/test_remote.py | 5 +++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py index d5fe013e0c..fc3284ca10 100644 --- a/tests/flytekit/unit/core/test_node_creation.py +++ b/tests/flytekit/unit/core/test_node_creation.py @@ -14,7 +14,6 @@ from flytekit.core.workflow import workflow from flytekit.exceptions.user import FlyteAssertion from flytekit.extras.accelerators import A100, T4 -from flytekit.image_spec.image_spec import ImageBuildEngine from flytekit.models import literals as _literal_models from flytekit.models.task import Resources as _resources_models from flytekit.tools.translator import get_serializable diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index d6b0cc711c..a048339a3e 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -472,6 +472,9 @@ def test_fetch_workflow_with_nested_branch(mock_promote, mock_workflow, remote): mock_promote.assert_called_with(ANY, node_launch_plans) +image_spec = ImageSpec(requirements="requirements.txt", registry="flyteorg") + + @mock.patch("pathlib.Path.read_bytes") @mock.patch("flytekit.remote.remote.FlyteRemote._version_from_hash") @mock.patch("flytekit.remote.remote.FlyteRemote.register_workflow") @@ -485,8 +488,6 @@ def test_get_image_names( compress_scripts_mock.return_value = "compressed" upload_file_mock.return_value = md5_bytes, "localhost:30084" - image_spec = ImageSpec(requirements="requirements.txt", registry="flyteorg") - @task(container_image=image_spec) def say_hello(name: str) -> str: return f"hello {name}!" From 9d4ff291282b55bf9888910d35dca27460a5389e Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 18:03:01 -0700 Subject: [PATCH 19/23] special hash Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 4 ++-- flytekit/image_spec/image_spec.py | 17 ++++++++++++++++- flytekit/tools/translator.py | 3 ++- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 60a6522ce8..1d795120c3 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -16,7 +16,7 @@ from flytekit.core.tracker import TrackedInstance, extract_task_module from flytekit.core.utils import _get_container_definition, _serialize_pod_spec, timeit from flytekit.extras.accelerators import BaseAccelerator -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduced_hash_from_image_spec from flytekit.loggers import logger from flytekit.models import task as _task_model from flytekit.models.security import Secret, SecurityContext @@ -276,7 +276,7 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :return: """ if isinstance(img, ImageSpec): - image = cfg.find_image(f"ft_{img.lhs}") + image = cfg.find_image(_calculate_deduced_hash_from_image_spec(img)) image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 83d67ae857..da52e63d13 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -23,7 +23,7 @@ @dataclass -class ImageSpec(TrackedInstance): +class ImageSpec: """ This class is used to specify the docker image that will be used to run the task. @@ -284,6 +284,21 @@ def _build_image(cls, builder, image_spec, img_name): cls._IMAGE_NAME_TO_REAL_NAME[img_name] = fully_qualified_image_name +@lru_cache +def _calculate_deduced_hash_from_image_spec(image_spec: ImageSpec): + """ + Calculate the hash from the image spec, + and it used to identify the imageSpec in the ImageConfig in the serialization context. + + ImageConfig: + - deduced hash 1: flyteorg/flytekit: 123 + - deduced hash 2: flyteorg/flytekit: 456 + """ + image_spec_bytes = asdict(image_spec).__str__().encode("utf-8") + # copy the image spec to avoid modifying the original image spec. otherwise, the hash will be different. + return base64.urlsafe_b64encode(hashlib.md5(image_spec_bytes).digest()).decode("ascii").rstrip("=") + + @lru_cache def calculate_hash_from_image_spec(image_spec: ImageSpec): """ diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 32dee5c5d9..6e8cd247c9 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -23,6 +23,7 @@ from flytekit.core.task import ReferenceTask from flytekit.core.utils import ClassDecorator, _dnsify from flytekit.core.workflow import ReferenceWorkflow, WorkflowBase +from flytekit.image_spec.image_spec import _calculate_deduced_hash_from_image_spec from flytekit.models import common as _common_models from flytekit.models import common as common_models from flytekit.models import interface as interface_models @@ -185,7 +186,7 @@ def get_serializable_task( if settings.image_config.images is None: settings.image_config = ImageConfig.create_from(settings.image_config.default_image) settings.image_config.images.append( - Image.look_up_image_info(f"ft_{e.container_image.lhs}", e.get_image(settings)) + Image.look_up_image_info(_calculate_deduced_hash_from_image_spec(e.container_image), e.get_image(settings)) ) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state From 06ee06f1821fc2d09b35d3609226a53658f92a2b Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 18:16:08 -0700 Subject: [PATCH 20/23] fix tests Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 6 +--- flytekit/tools/translator.py | 4 ++- .../flytekit/unit/core/test_container_task.py | 33 +++++++++---------- tests/flytekit/unit/core/test_dynamic.py | 2 +- .../unit/core/test_python_auto_container.py | 9 +++-- .../unit/core/test_python_function_task.py | 4 +-- .../flytekit/unit/core/test_serialization.py | 19 ++++++----- tests/flytekit/unit/remote/test_remote.py | 5 ++- 8 files changed, 38 insertions(+), 44 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index da52e63d13..fec45a813a 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -14,7 +14,6 @@ import requests from packaging.version import Version -from flytekit.core.tracker import TrackedInstance from flytekit.exceptions.user import FlyteAssertion DOCKER_HUB = "docker.io" @@ -73,9 +72,6 @@ class ImageSpec: tag_format: Optional[str] = None def __post_init__(self): - self._instantiated_in = None - self._module_file = None - self._lhs = None self.name = self.name.lower() self._is_force_push = os.environ.get(FLYTE_FORCE_PUSH_IMAGE_SPEC, False) # False by default if self.registry: @@ -287,7 +283,7 @@ def _build_image(cls, builder, image_spec, img_name): @lru_cache def _calculate_deduced_hash_from_image_spec(image_spec: ImageSpec): """ - Calculate the hash from the image spec, + Calculate this special hash from the image spec, and it used to identify the imageSpec in the ImageConfig in the serialization context. ImageConfig: diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 6e8cd247c9..04a0eede63 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -186,7 +186,9 @@ def get_serializable_task( if settings.image_config.images is None: settings.image_config = ImageConfig.create_from(settings.image_config.default_image) settings.image_config.images.append( - Image.look_up_image_info(_calculate_deduced_hash_from_image_spec(e.container_image), e.get_image(settings)) + Image.look_up_image_info( + _calculate_deduced_hash_from_image_spec(e.container_image), e.get_image(settings) + ) ) # In case of Dynamic tasks, we want to pass the serialization context, so that they can reconstruct the state diff --git a/tests/flytekit/unit/core/test_container_task.py b/tests/flytekit/unit/core/test_container_task.py index 66d0ef1999..1281a9ec14 100644 --- a/tests/flytekit/unit/core/test_container_task.py +++ b/tests/flytekit/unit/core/test_container_task.py @@ -158,11 +158,9 @@ def test_pod_template(): assert serialized_pod_spec["runtimeClassName"] == "nvidia" -image_spec = ImageSpec(registry="flyte", base_image="r-base", builder="test-raw-container") - - def test_raw_container_with_image_spec(mock_image_spec_builder): ImageBuildEngine.register("test-raw-container", mock_image_spec_builder) + image_spec = ImageSpec(registry="flyte", base_image="r-base", builder="test-raw-container") calculate_ellipse_area_r = ContainerTask( name="ellipse-area-metadata-r", @@ -188,21 +186,6 @@ def test_raw_container_with_image_spec(mock_image_spec_builder): assert container.image == image_spec.image_name() -image_spec_1 = ImageSpec( - name="image-1", - packages=["numpy"], - registry="localhost:30000", - builder="test", -) - -image_spec_2 = ImageSpec( - name="image-2", - packages=["pandas"], - registry="localhost:30000", - builder="test", -) - - def test_container_task_image_spec(mock_image_spec_builder): default_image = Image(name="default", fqn="docker.io/xyz", tag="some-git-hash") default_image_config = ImageConfig(default_image=default_image) @@ -211,6 +194,20 @@ def test_container_task_image_spec(mock_image_spec_builder): project="p", domain="d", version="v", image_config=default_image_config, env={"FOO": "bar"} ) + image_spec_1 = ImageSpec( + name="image-1", + packages=["numpy"], + registry="localhost:30000", + builder="test", + ) + + image_spec_2 = ImageSpec( + name="image-2", + packages=["pandas"], + registry="localhost:30000", + builder="test", + ) + ps = V1PodSpec( containers=[ V1Container( diff --git a/tests/flytekit/unit/core/test_dynamic.py b/tests/flytekit/unit/core/test_dynamic.py index e49e09fad6..7964548674 100644 --- a/tests/flytekit/unit/core/test_dynamic.py +++ b/tests/flytekit/unit/core/test_dynamic.py @@ -287,7 +287,7 @@ def dt(mode: int) -> int: raise ValueError("Invalid mode") entity_mapping = OrderedDict() - get_serializable_task(entity_mapping, settings.with_serialized_context(), dt) + get_serializable_task(entity_mapping, settings, dt) serialised_entities_iterator = iter(entity_mapping.values()) assert "t1" in next(serialised_entities_iterator).template.id.name diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index 142e4f86f4..5db1d717d5 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -9,7 +9,7 @@ from flytekit.core.pod_template import PodTemplate from flytekit.core.python_auto_container import PythonAutoContainerTask, get_registerable_container_image from flytekit.core.resources import Resources -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduced_hash_from_image_spec from flytekit.tools.translator import get_serializable_task @@ -54,13 +54,12 @@ def serialization_settings(request): return request.getfixturevalue(request.param) -image_spec = ImageSpec(name="image-1", registry="localhost:30000", builder="test") - - def test_image_name_interpolation(default_image_config): + image_spec = ImageSpec(name="image-1", registry="localhost:30000", builder="test") + new_img_cfg = ImageConfig.create_from( default_image_config.default_image, - other_images=[Image.look_up_image_info(f"ft_{image_spec.lhs}", "flyte/test:d1")], + other_images=[Image.look_up_image_info(_calculate_deduced_hash_from_image_spec(image_spec), "flyte/test:d1")], ) img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" img = get_registerable_container_image(img=img_to_interpolate, cfg=new_img_cfg) diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index a268ecb51c..8ba875b664 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -36,9 +36,6 @@ def test_istestfunction(): assert istestfunction(tasks.tasks) is False -image_spec = ImageSpec(builder="test", python_version="3.7", registry="") - - def test_container_image_conversion(mock_image_spec_builder): default_img = Image(name="default", fqn="xyz.com/abc", tag="tag1") other_img = Image(name="other", fqn="xyz.com/other", tag="tag-other") @@ -114,6 +111,7 @@ def test_container_image_conversion(mock_image_spec_builder): ) ImageBuildEngine.register("test", mock_image_spec_builder) + image_spec = ImageSpec(builder="test", python_version="3.7", registry="") assert get_registerable_container_image(image_spec, cfg) == image_spec.image_name() diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index 9805aea56c..bced85df52 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -12,7 +12,7 @@ from flytekit.core.python_auto_container import get_registerable_container_image from flytekit.core.task import task from flytekit.core.workflow import workflow -from flytekit.image_spec.image_spec import ImageBuildEngine +from flytekit.image_spec.image_spec import ImageBuildEngine, _calculate_deduced_hash_from_image_spec from flytekit.models.admin.workflow import WorkflowSpec from flytekit.models.types import SimpleType from flytekit.tools.translator import get_serializable @@ -26,12 +26,6 @@ env=None, image_config=ImageConfig(default_image=default_img, images=[default_img]), ) -image_spec = ImageSpec( - packages=["mypy"], - apt_packages=["git"], - registry="ghcr.io/flyteorg", - builder="test", -) def test_serialization(): @@ -280,6 +274,13 @@ def t5(a: int) -> int: def t6(a: int) -> int: return a + image_spec = ImageSpec( + packages=["mypy"], + apt_packages=["git"], + registry="ghcr.io/flyteorg", + builder="test", + ) + @task(container_image=image_spec) def t7(a: int) -> int: return a @@ -288,7 +289,9 @@ def t7(a: int) -> int: imgs = ImageConfig.auto( config_file=os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/images.config") ) - imgs.images.append(Image(name=f"ft_{image_spec.lhs}", fqn="docker.io/t7", tag="latest")) + imgs.images.append( + Image(name=_calculate_deduced_hash_from_image_spec(image_spec), fqn="docker.io/t7", tag="latest") + ) rs = flytekit.configuration.SerializationSettings( project="project", domain="domain", diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index a048339a3e..d6b0cc711c 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -472,9 +472,6 @@ def test_fetch_workflow_with_nested_branch(mock_promote, mock_workflow, remote): mock_promote.assert_called_with(ANY, node_launch_plans) -image_spec = ImageSpec(requirements="requirements.txt", registry="flyteorg") - - @mock.patch("pathlib.Path.read_bytes") @mock.patch("flytekit.remote.remote.FlyteRemote._version_from_hash") @mock.patch("flytekit.remote.remote.FlyteRemote.register_workflow") @@ -488,6 +485,8 @@ def test_get_image_names( compress_scripts_mock.return_value = "compressed" upload_file_mock.return_value = md5_bytes, "localhost:30084" + image_spec = ImageSpec(requirements="requirements.txt", registry="flyteorg") + @task(container_image=image_spec) def say_hello(name: str) -> str: return f"hello {name}!" From a8f455735c83dbf3079c39a4381a93d0f951c800 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 18:21:10 -0700 Subject: [PATCH 21/23] fix tests Signed-off-by: Kevin Su --- tests/flytekit/unit/core/test_node_creation.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py index fc3284ca10..684f49031b 100644 --- a/tests/flytekit/unit/core/test_node_creation.py +++ b/tests/flytekit/unit/core/test_node_creation.py @@ -14,12 +14,15 @@ from flytekit.core.workflow import workflow from flytekit.exceptions.user import FlyteAssertion from flytekit.extras.accelerators import A100, T4 +from flytekit.image_spec.image_spec import ImageBuildEngine from flytekit.models import literals as _literal_models from flytekit.models.task import Resources as _resources_models from flytekit.tools.translator import get_serializable -def test_normal_task(): +def test_normal_task(mock_image_spec_builder): + ImageBuildEngine.register("test", mock_image_spec_builder) + @task def t1(a: str) -> str: return a + " world" From 7584df3677b56834044396eb2053cf51409e9ce9 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 18:23:50 -0700 Subject: [PATCH 22/23] nit Signed-off-by: Kevin Su --- flytekit/core/python_auto_container.py | 4 ++-- flytekit/image_spec/image_spec.py | 2 +- flytekit/tools/translator.py | 4 ++-- tests/flytekit/unit/core/test_python_auto_container.py | 4 ++-- tests/flytekit/unit/core/test_serialization.py | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) diff --git a/flytekit/core/python_auto_container.py b/flytekit/core/python_auto_container.py index 1d795120c3..2c4703cdd3 100644 --- a/flytekit/core/python_auto_container.py +++ b/flytekit/core/python_auto_container.py @@ -16,7 +16,7 @@ from flytekit.core.tracker import TrackedInstance, extract_task_module from flytekit.core.utils import _get_container_definition, _serialize_pod_spec, timeit from flytekit.extras.accelerators import BaseAccelerator -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduced_hash_from_image_spec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduped_hash_from_image_spec from flytekit.loggers import logger from flytekit.models import task as _task_model from flytekit.models.security import Secret, SecurityContext @@ -276,7 +276,7 @@ def get_registerable_container_image(img: Optional[Union[str, ImageSpec]], cfg: :return: """ if isinstance(img, ImageSpec): - image = cfg.find_image(_calculate_deduced_hash_from_image_spec(img)) + image = cfg.find_image(_calculate_deduped_hash_from_image_spec(img)) image_name = image.full if image else None if not image_name: ImageBuildEngine.build(img) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index fec45a813a..37f87549d0 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -281,7 +281,7 @@ def _build_image(cls, builder, image_spec, img_name): @lru_cache -def _calculate_deduced_hash_from_image_spec(image_spec: ImageSpec): +def _calculate_deduped_hash_from_image_spec(image_spec: ImageSpec): """ Calculate this special hash from the image spec, and it used to identify the imageSpec in the ImageConfig in the serialization context. diff --git a/flytekit/tools/translator.py b/flytekit/tools/translator.py index 04a0eede63..a77e0a0bf5 100644 --- a/flytekit/tools/translator.py +++ b/flytekit/tools/translator.py @@ -23,7 +23,7 @@ from flytekit.core.task import ReferenceTask from flytekit.core.utils import ClassDecorator, _dnsify from flytekit.core.workflow import ReferenceWorkflow, WorkflowBase -from flytekit.image_spec.image_spec import _calculate_deduced_hash_from_image_spec +from flytekit.image_spec.image_spec import _calculate_deduped_hash_from_image_spec from flytekit.models import common as _common_models from flytekit.models import common as common_models from flytekit.models import interface as interface_models @@ -187,7 +187,7 @@ def get_serializable_task( settings.image_config = ImageConfig.create_from(settings.image_config.default_image) settings.image_config.images.append( Image.look_up_image_info( - _calculate_deduced_hash_from_image_spec(e.container_image), e.get_image(settings) + _calculate_deduped_hash_from_image_spec(e.container_image), e.get_image(settings) ) ) diff --git a/tests/flytekit/unit/core/test_python_auto_container.py b/tests/flytekit/unit/core/test_python_auto_container.py index 5db1d717d5..5068da53de 100644 --- a/tests/flytekit/unit/core/test_python_auto_container.py +++ b/tests/flytekit/unit/core/test_python_auto_container.py @@ -9,7 +9,7 @@ from flytekit.core.pod_template import PodTemplate from flytekit.core.python_auto_container import PythonAutoContainerTask, get_registerable_container_image from flytekit.core.resources import Resources -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduced_hash_from_image_spec +from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, _calculate_deduped_hash_from_image_spec from flytekit.tools.translator import get_serializable_task @@ -59,7 +59,7 @@ def test_image_name_interpolation(default_image_config): new_img_cfg = ImageConfig.create_from( default_image_config.default_image, - other_images=[Image.look_up_image_info(_calculate_deduced_hash_from_image_spec(image_spec), "flyte/test:d1")], + other_images=[Image.look_up_image_info(_calculate_deduped_hash_from_image_spec(image_spec), "flyte/test:d1")], ) img_to_interpolate = "{{.image.default.fqn}}:{{.image.default.version}}-special" img = get_registerable_container_image(img=img_to_interpolate, cfg=new_img_cfg) diff --git a/tests/flytekit/unit/core/test_serialization.py b/tests/flytekit/unit/core/test_serialization.py index bced85df52..88297f43f4 100644 --- a/tests/flytekit/unit/core/test_serialization.py +++ b/tests/flytekit/unit/core/test_serialization.py @@ -12,7 +12,7 @@ from flytekit.core.python_auto_container import get_registerable_container_image from flytekit.core.task import task from flytekit.core.workflow import workflow -from flytekit.image_spec.image_spec import ImageBuildEngine, _calculate_deduced_hash_from_image_spec +from flytekit.image_spec.image_spec import ImageBuildEngine, _calculate_deduped_hash_from_image_spec from flytekit.models.admin.workflow import WorkflowSpec from flytekit.models.types import SimpleType from flytekit.tools.translator import get_serializable @@ -290,7 +290,7 @@ def t7(a: int) -> int: config_file=os.path.join(os.path.dirname(os.path.realpath(__file__)), "configs/images.config") ) imgs.images.append( - Image(name=_calculate_deduced_hash_from_image_spec(image_spec), fqn="docker.io/t7", tag="latest") + Image(name=_calculate_deduped_hash_from_image_spec(image_spec), fqn="docker.io/t7", tag="latest") ) rs = flytekit.configuration.SerializationSettings( project="project", From fb0366271fba507bb20c973fa0c6272bc686334c Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Mon, 3 Jun 2024 18:44:06 -0700 Subject: [PATCH 23/23] nit Signed-off-by: Kevin Su --- plugins/flytekit-envd/flytekitplugins/envd/image_builder.py | 2 +- plugins/flytekit-envd/tests/test_image_spec.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 59b8ca07eb..344003bafd 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -91,7 +91,7 @@ def create_envd_config(image_spec: ImageSpec) -> str: run_commands = _create_str_from_package_list(image_spec.commands) conda_channels = _create_str_from_package_list(image_spec.conda_channels) apt_packages = _create_str_from_package_list(image_spec.apt_packages) - env = {"PYTHONPATH": "/root", _F_IMG_ID: image_spec.image_name()} + env = {"PYTHONPATH": "/root:", _F_IMG_ID: image_spec.image_name()} if image_spec.env: env.update(image_spec.env) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 7fd3cd1be0..5b7b73f755 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -57,7 +57,7 @@ def build(): run(commands=["echo hello"]) install.python_packages(name=["pandas"]) install.apt_packages(name=["git"]) - runtime.environ(env={{'PYTHONPATH': '/root', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) + runtime.environ(env={{'PYTHONPATH': '/root:', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) config.pip_index(url="https://private-pip-index/simple") install.python(version="3.8") io.copy(source="./", target="/root") @@ -88,7 +88,7 @@ def build(): run(commands=[]) install.python_packages(name=["flytekit"]) install.apt_packages(name=[]) - runtime.environ(env={{'PYTHONPATH': '/root', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) + runtime.environ(env={{'PYTHONPATH': '/root:', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) config.pip_index(url="https://pypi.org/simple") install.conda(use_mamba=True) install.conda_packages(name=["pytorch", "cpuonly"], channel=["pytorch"]) @@ -122,7 +122,7 @@ def build(): run(commands=[]) install.python_packages(name=["-U --pre pandas", "torch", "torchvision"]) install.apt_packages(name=[]) - runtime.environ(env={{'PYTHONPATH': '/root', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) + runtime.environ(env={{'PYTHONPATH': '/root:', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root']) config.pip_index(url="https://pypi.org/simple", extra_url="https://download.pytorch.org/whl/cpu https://pypi.anaconda.org/scientific-python-nightly-wheels/simple") """ )