From 93179b14e4182565c65f0ce1abb60148eda9fcbb Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sun, 30 Apr 2023 14:25:58 -0700 Subject: [PATCH] Add is_inside func to ImageSpec (#1601) * Add is_inside func to ImageSpec Signed-off-by: Kevin Su * test Signed-off-by: Kevin Su * Add is_inside func to ImageSpec Signed-off-by: Kevin Su * Add tests Signed-off-by: Kevin Su * rename Signed-off-by: Kevin Su * rename Signed-off-by: Kevin Su * check execution mode Signed-off-by: Kevin Su --------- Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 12 +++++++++++- .../flytekitplugins/envd/image_builder.py | 4 ++-- plugins/flytekit-envd/tests/test_image_spec.py | 2 +- .../unit/core/image_spec/test_image_spec.py | 14 +++++++++++++- 4 files changed, 27 insertions(+), 5 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 72c815c574..3ba118741b 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -13,6 +13,8 @@ from dataclasses_json import dataclass_json from docker.errors import APIError, ImageNotFound +_F_IMG_ID = "_F_IMG_ID" + @dataclass_json @dataclass @@ -23,9 +25,9 @@ class ImageSpec: Args: name: name of the image. python_version: python version of the image. + builder: Type of plugin to build the image. Use envd by default. source_root: source root of the image. env: environment variables of the image. - destination_dir: destination directory of the image. registry: registry of the image. packages: list of python packages to install. apt_packages: list of apt packages to install. @@ -52,6 +54,14 @@ def image_name(self) -> str: container_image = f"{self.registry}/{container_image}" return container_image + def is_container(self) -> bool: + from flytekit.core.context_manager import ExecutionState, FlyteContextManager + + state = FlyteContextManager.current_context().execution_state + if state and state.mode and state.mode != ExecutionState.Mode.LOCAL_WORKFLOW_EXECUTION: + return os.environ.get(_F_IMG_ID) == self.image_name() + return True + def exist(self) -> bool: """ Check if the image exists in the registry. diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 22e46c787f..e5426b8731 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -6,7 +6,7 @@ from flytekit.configuration import DefaultImages from flytekit.core import context_manager -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec, ImageSpecBuilder +from flytekit.image_spec.image_spec import _F_IMG_ID, ImageBuildEngine, ImageSpec, ImageSpecBuilder class EnvdImageSpecBuilder(ImageSpecBuilder): @@ -39,7 +39,7 @@ def create_envd_config(image_spec: ImageSpec) -> str: packages = [] if image_spec.packages is None else image_spec.packages apt_packages = [] if image_spec.apt_packages is None else image_spec.apt_packages env = {} if image_spec.env is None else image_spec.env - env.update({"PYTHONPATH": "/root"}) + env.update({"PYTHONPATH": "/root", _F_IMG_ID: image_spec.image_name()}) envd_config = f"""# syntax=v1 diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 96c06c93ba..87b0668171 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -26,6 +26,6 @@ def build(): install.python_packages(name = ["pandas"]) install.apt_packages(name = ["git"]) install.python(version="3.8") - runtime.environ(env={'PYTHONPATH': '/root'}) + runtime.environ(env={'PYTHONPATH': '/root', '_F_IMG_ID': 'flytekit:yZ8jICcDTLoDArmNHbWNwg..'}) """ ) 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 d2d4975364..8a1cfc75f6 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -1,7 +1,11 @@ +import os + import pytest +from flytekit.core import context_manager +from flytekit.core.context_manager import ExecutionState from flytekit.image_spec import ImageSpec -from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpecBuilder, calculate_hash_from_image_spec +from flytekit.image_spec.image_spec import _F_IMG_ID, ImageBuildEngine, ImageSpecBuilder, calculate_hash_from_image_spec def test_image_spec(): @@ -22,6 +26,14 @@ def test_image_spec(): assert image_spec.builder == "envd" assert image_spec.source_root is None assert image_spec.env is None + assert image_spec.is_container() is True + assert image_spec.image_name() == "flytekit:yZ8jICcDTLoDArmNHbWNwg.." + ctx = context_manager.FlyteContext.current_context() + with context_manager.FlyteContextManager.with_context( + ctx.with_execution_state(ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION)) + ): + os.environ[_F_IMG_ID] = "flytekit:123" + assert image_spec.is_container() is False class DummyImageSpecBuilder(ImageSpecBuilder): def build_image(self, img):