Add support overriding image (#1652)

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · May 24, 2023 · eafcc82 · eafcc82
1 parent 67b92be
commit eafcc82
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/flytekit/core/node.py b/flytekit/core/node.py
@@ -128,6 +128,8 @@ def with_overrides(self, *args, **kwargs):
             if not isinstance(new_task_config, type(self.flyte_entity._task_config)):
                 raise ValueError("can't change the type of the task config")
             self.flyte_entity._task_config = new_task_config
+        if "container_image" in kwargs:
+            self.flyte_entity._container_image = kwargs["container_image"]
         return self
 
 

diff --git a/tests/flytekit/unit/core/test_node_creation.py b/tests/flytekit/unit/core/test_node_creation.py
@@ -452,3 +452,16 @@ def my_wf(a: str) -> str:
             return t1(a=a).with_overrides(task_config=None)
 
         my_wf()
+
+
+def test_override_image():
+    @task
+    def bar():
+        print("hello")
+
+    @workflow
+    def wf() -> str:
+        bar().with_overrides(container_image="hello/world")
+        return "hi"
+
+    assert wf.nodes[0].flyte_entity.container_image == "hello/world"