Local execution with compiled graph(#4080)

flytekit/core/condition.py

@@ -18,19 +18,46 @@
 )
 from flytekit.models.core import condition as _core_cond
 from flytekit.models.core import workflow as _core_wf
+from flytekit.models.core.workflow import IfElseBlock
 from flytekit.models.literals import Binding, BindingData, Literal, RetryStrategy
 from flytekit.models.types import Error
 
 
 class BranchNode(object):
-    def __init__(self, name: str, ifelse_block: _core_wf.IfElseBlock):
+    def __init__(self, name: str, ifelse_block: _core_wf.IfElseBlock, cs: typing.Optional[ConditionalSection] = None):
         self._name = name
         self._ifelse_block = ifelse_block
+        self._cs = cs
 
     @property
     def name(self):
         return self._name
 
+    # Output Node or None
+    def __call__(self, **kwargs):
+        for c in self._cs.cases:
+            c._expr = _update_promise(c.expr, kwargs)
+            if c.expr is None:
+                return c.output_node
+            if c.expr.eval():
+                return c.output_node
+
+
+def _update_promise(operand: Union[Literal, Promise, ConjunctionExpression, ComparisonExpression], promises: typing.Dict[str, Promise]):
+    if isinstance(operand, Literal):
+        return Promise(var="placeholder", val=operand)
+    if isinstance(operand, ConjunctionExpression):
+        lhs = _update_promise(operand.lhs, promises)
+        rhs = _update_promise(operand.rhs, promises)
+        op = operand.op
+        return ConjunctionExpression(lhs, op, rhs)  # type: ignore
+    if isinstance(operand, ComparisonExpression):
+        lhs = _update_promise(operand.lhs, promises)
+        rhs = _update_promise(operand.rhs, promises)
+        op = operand.op
+        return ComparisonExpression(lhs, op, rhs)
+    return promises[create_branch_node_promise_var(operand.ref.node_id, operand.var)]
+
 
 class ConditionalSection:
     """
@@ -84,6 +111,7 @@ def end_branch(self) -> Optional[Union[Condition, Promise, Tuple[Promise], VoidP
         If so then return the promise, else return the condition
         """
         if self._last_case:
+            print("last")
             # We have completed the conditional section, lets pop off the branch context
             FlyteContextManager.pop_context()
             ctx = FlyteContextManager.current_context()
@@ -438,13 +466,13 @@ def transform_to_boolexpr(
 
 
 def to_case_block(c: Case) -> Tuple[Union[_core_wf.IfBlock], typing.List[Promise]]:
+    if c.output_promise is None:
+        raise AssertionError("Illegal Condition block, with no output promise")
     expr, promises = transform_to_boolexpr(cast(Union[ComparisonExpression, ConjunctionExpression], c.expr))
-    if c.output_promise is not None:
-        n = c.output_node
-    return _core_wf.IfBlock(condition=expr, then_node=n), promises
+    return _core_wf.IfBlock(condition=expr, then_node=c.output_node), promises
 
 
-def to_ifelse_block(node_id: str, cs: ConditionalSection) -> Tuple[_core_wf.IfElseBlock, typing.List[Binding]]:
+def to_ifelse_block(node_id: str, cs: ConditionalSection) -> tuple[IfElseBlock, list[Promise]]:
     if len(cs.cases) == 0:
         raise AssertionError("Illegal Condition block, with no if-else cases")
     if len(cs.cases) < 2:
@@ -474,7 +502,7 @@ def to_ifelse_block(node_id: str, cs: ConditionalSection) -> Tuple[_core_wf.IfEl
 
 def to_branch_node(name: str, cs: ConditionalSection) -> Tuple[BranchNode, typing.List[Promise]]:
     blocks, promises = to_ifelse_block(name, cs)
-    return BranchNode(name=name, ifelse_block=blocks), promises
+    return BranchNode(name=name, ifelse_block=blocks, cs=cs), promises
 
 
 def conditional(name: str) -> ConditionalSection:

flytekit/core/gate.py

@@ -111,9 +111,12 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr
             return p
 
         # Assume this is an approval operation since that's the only remaining option.
-        msg = click.style("[Approval Gate] ", fg="yellow") + click.style(
-            f"@{self.name} Approve {typing.cast(Promise, self._upstream_item).val.value}?", fg="cyan"
-        )
+        assert len(kwargs) == 1
+
+        value = kwargs[list(kwargs.keys())[0]]
+        if isinstance(value, Promise):
+            value = value.eval()
+        msg = click.style("[Approval Gate] ", fg="yellow") + click.style(f"@{self.name} Approve {value}?", fg="cyan")
         proceed = click.confirm(msg, default=True)
         if proceed:
             # We need to return a promise here, and a promise is what should've been passed in by the call in approve()
@@ -127,6 +130,9 @@ def local_execute(self, ctx: FlyteContext, **kwargs) -> Union[Tuple[Promise], Pr
     def local_execution_mode(self):
         return ExecutionState.Mode.LOCAL_TASK_EXECUTION
 
+    def __call__(self, *args: object, **kwargs: object) -> Union[Tuple[Promise], Promise, VoidPromise, Tuple, None]:
+        return flyte_entity_call_handler(self, *args, **kwargs)  # type: ignore
+
 
 def wait_for_input(name: str, timeout: datetime.timedelta, expected_type: typing.Type):
     """Create a Gate object that waits for user input of the specified type.

flytekit/core/promise.py

@@ -299,7 +299,7 @@ def t1() -> (int, str): ...
 
     # TODO: Currently, NodeOutput we're creating is the slimmer core package Node class, but since only the
     #  id is used, it's okay for now. Let's clean all this up though.
-    def __init__(self, var: str, val: Union[NodeOutput, _literals_models.Literal]):
+    def __init__(self, var: str, val: Optional[Union[NodeOutput, _literals_models.Literal]]):
         self._var = var
         self._promise_ready = True
         self._val = val
@@ -1043,6 +1043,7 @@ def flyte_entity_call_handler(
                     return None
             return cast(LocallyExecutable, entity).local_execute(ctx, **kwargs)
     else:
+        # If LocallyExecutable is workflow, then we would get into here,
         mode = cast(LocallyExecutable, entity).local_execution_mode()
         with FlyteContextManager.with_context(
             ctx.with_execution_state(ctx.new_execution_state().with_params(mode=mode))

flytekit/core/python_function_task.py

@@ -281,8 +281,20 @@ def dynamic_execute(self, task_function: Callable, **kwargs) -> Any:
             # The rest of this function mimics the local_execute of the workflow. We can't use the workflow
             # local_execute directly though since that converts inputs into Promises.
             logger.debug(f"Executing Dynamic workflow, using raw inputs {kwargs}")
-            self._create_and_cache_dynamic_workflow()
-            function_outputs = cast(PythonFunctionWorkflow, self._wf).execute(**kwargs)
+            if not ctx.compilation_state:
+                cs = ctx.new_compilation_state(prefix="d")
+            else:
+                cs = ctx.compilation_state.with_params(prefix="d")
+
+            updated_ctx = ctx.with_compilation_state(cs)
+            if self.execution_mode == self.ExecutionBehavior.DYNAMIC:
+                es = ctx.new_execution_state().with_params(mode=ExecutionState.Mode.DYNAMIC_TASK_EXECUTION)
+                updated_ctx = updated_ctx.with_execution_state(es)
+
+            with FlyteContextManager.with_context(updated_ctx):
+                self._create_and_cache_dynamic_workflow()
+                cast(PythonFunctionWorkflow, self._wf).compile(**kwargs)
+            function_outputs = self._wf.execute(**kwargs)
 
             if isinstance(function_outputs, VoidPromise) or function_outputs is None:
                 return VoidPromise(self.name)