Add is_inside func to ImageSpec (#1601)

* Add is_inside func to ImageSpec Signed-off-by: Kevin Su <pingsutw@apache.org> * test Signed-off-by: Kevin Su <pingsutw@apache.org> * Add is_inside func to ImageSpec Signed-off-by: Kevin Su <pingsutw@apache.org> * Add tests Signed-off-by: Kevin Su <pingsutw@apache.org> * rename Signed-off-by: Kevin Su <pingsutw@apache.org> * rename Signed-off-by: Kevin Su <pingsutw@apache.org> * check execution mode Signed-off-by: Kevin Su <pingsutw@apache.org> --------- Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Apr 30, 2023 · 4ad64d3 · 4ad64d3
1 parent 36fc151
commit 4ad64d3
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):