Use imageSpec name as key

Signed-off-by: Kevin Su <[email protected]>

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)

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")

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