Add is_inside func to ImageSpec (#1601)

* Add is_inside func to ImageSpec Signed-off-by: Kevin Su <[email protected]> * test Signed-off-by: Kevin Su <[email protected]> * Add is_inside func to ImageSpec Signed-off-by: Kevin Su <[email protected]> * Add tests Signed-off-by: Kevin Su <[email protected]> * rename Signed-off-by: Kevin Su <[email protected]> * rename Signed-off-by: Kevin Su <[email protected]> * check execution mode Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]>
flyteorg · May 16, 2023 · 93179b1 · 93179b1
1 parent 21bd0a7
commit 93179b1
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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):