Add is_inside func to ImageSpec

flytekit/image_spec/image_spec.py

@@ -13,6 +13,8 @@
 from dataclasses_json import dataclass_json
 from docker.errors import APIError, ImageNotFound
 
+FLYTE_IMAGE_NAME = "FLYTE_IMAGE_NAME"
+
 
 @dataclass_json
 @dataclass
@@ -52,6 +54,11 @@ def image_name(self) -> str:
             container_image = f"{self.registry}/{container_image}"
         return container_image
 
+    def is_inside(self) -> bool:
+        if os.environ.get(FLYTE_IMAGE_NAME):
+            return os.environ.get(FLYTE_IMAGE_NAME) == self.image_name()
+        return True
+
     def exist(self) -> bool:
         """
         Check if the image exists in the registry.

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 FLYTE_IMAGE_NAME, 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", FLYTE_IMAGE_NAME: image_spec.image_name()})
 
     envd_config = f"""# syntax=v1
 

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -1,7 +1,14 @@
+import os
+
 import pytest
 
 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 (
+    FLYTE_IMAGE_NAME,
+    ImageBuildEngine,
+    ImageSpecBuilder,
+    calculate_hash_from_image_spec,
+)
 
 
 def test_image_spec():
@@ -22,6 +29,10 @@ 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_inside() is True
+    assert image_spec.image_name() == "flytekit:yZ8jICcDTLoDArmNHbWNwg.."
+    os.environ[FLYTE_IMAGE_NAME] = "flytekit:123"
+    assert image_spec.is_inside() is False
 
     class DummyImageSpecBuilder(ImageSpecBuilder):
         def build_image(self, img):