Fetch task executions in dynamic (#1636)

* fetch task executions in dynamic Signed-off-by: Kevin Su <[email protected]> * nit Signed-off-by: Kevin Su <[email protected]> * lint Signed-off-by: Kevin Su <[email protected]> --------- Signed-off-by: Kevin Su <[email protected]>
flyteorg · May 15, 2023 · faf02b1 · faf02b1
1 parent b410108
commit faf02b1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/flytekit/remote/entities.py b/flytekit/remote/entities.py
@@ -362,6 +362,7 @@ def __init__(
             )
         # TODO: Revisit flyte_branch_node and flyte_gate_node, should they be another type like Condition instead
         #       of a node?
+        self._flyte_task_node = task_node
         if task_node:
             self._flyte_entity = task_node.flyte_task
         elif workflow_node:
@@ -382,6 +383,10 @@ def __init__(
         )
         self._upstream = upstream_nodes
 
+    @property
+    def task_node(self) -> Optional[FlyteTaskNode]:
+        return self._flyte_task_node
+
     @property
     def flyte_entity(self) -> Union[FlyteTask, FlyteWorkflow, FlyteLaunchPlan, FlyteBranchNode]:
         return self._flyte_entity
@@ -707,7 +712,7 @@ def promote_from_closure(
 
         :param closure: This is the closure returned by Admin
         :param node_launch_plans: The reason this exists is because the compiled closure doesn't have launch plans.
-          It only has subworkflows and tasks. Why this is is unclear. If supplied, this map of launch plans will be
+          It only has subworkflows and tasks. Why this is unclear. If supplied, this map of launch plans will be
         :return:
         """
         sub_workflows = {sw.template.id: sw.template for sw in closure.sub_workflows}

diff --git a/flytekit/remote/remote.py b/flytekit/remote/remote.py
@@ -1632,13 +1632,10 @@ def sync_node_execution(
                     self.sync_node_execution(FlyteNodeExecution.promote_from_model(cne), dynamic_flyte_wf._node_map)
                     for cne in child_node_executions
                 ]
-                # This is copied from below - dynamic tasks have both task executions (executions of the parent
-                # task) as well as underlying node executions (of the generated subworkflow). Feel free to refactor
-                # if you can think of a better way.
                 execution._task_executions = [
-                    self.sync_task_execution(FlyteTaskExecution.promote_from_model(t))
-                    for t in iterate_task_executions(self.client, execution.id)
+                    node_exes.task_executions for node_exes in execution.subworkflow_node_executions.values()
                 ]
+
                 execution._interface = dynamic_flyte_wf.interface
 
             # Handle the case where it's a static subworkflow
@@ -1665,7 +1662,9 @@ def sync_node_execution(
         # This is the plain ol' task execution case
         else:
             execution._task_executions = [
-                self.sync_task_execution(FlyteTaskExecution.promote_from_model(t))
+                self.sync_task_execution(
+                    FlyteTaskExecution.promote_from_model(t), node_mapping[node_id].task_node.flyte_task
+                )
                 for t in iterate_task_executions(self.client, execution.id)
             ]
             execution._interface = execution._node.flyte_entity.interface
@@ -1679,17 +1678,15 @@ def sync_node_execution(
         return execution
 
     def sync_task_execution(
-        self, execution: FlyteTaskExecution, entity_definition: typing.Union[FlyteWorkflow, FlyteTask] = None
+        self, execution: FlyteTaskExecution, entity_definition: typing.Optional[FlyteTask] = None
     ) -> FlyteTaskExecution:
         """Sync a FlyteTaskExecution object with its corresponding remote state."""
-        if entity_definition is not None:
-            raise ValueError("Entity definition arguments aren't supported when syncing task executions")
-
         execution._closure = self.client.get_task_execution(execution.id).closure
         execution_data = self.client.get_task_execution_data(execution.id)
         task_id = execution.id.task_id
-        task = self.fetch_task(task_id.project, task_id.domain, task_id.name, task_id.version)
-        return self._assign_inputs_and_outputs(execution, execution_data, task.interface)
+        if entity_definition is None:
+            entity_definition = self.fetch_task(task_id.project, task_id.domain, task_id.name, task_id.version)
+        return self._assign_inputs_and_outputs(execution, execution_data, entity_definition.interface)
 
     #############################
     # Terminate Execution State #